2778. Sum of Squares of Special Elements

Array
Math
Enumeration

TIME

O(sqrt(n))

SPACE

O(1)

Intuition

The special positions are exactly the divisors of n, and there are far fewer of those than there are elements. Rather than testing n % i == 0 for every index, enumerate the divisors directly: walk i from 1 while i * i <= n, and whenever i divides n, both i and n / i are divisors. That touches only O(sqrt(n)) indices.

The one wrinkle is indexing. The statement is 1-indexed, so divisor i refers to nums[i - 1].

Approach

n = len(nums), total = 0.
For i from 1 while i * i <= n:
- if n % i != 0, continue
- add nums[i - 1] ** 2
- let j = n // i; if j != i, add nums[j - 1] ** 2
Return total.

The j != i guard stops a perfect square from being counted twice, since i and n / i coincide when i * i == n.

Complexity

Time: O(sqrt(n))
Space: O(1)

At n <= 50 the simple full scan is equally fast in practice; the divisor enumeration is used because it is the shape the problem actually has.

Edge Cases

n = 1 → only index 1 is special, so the answer is nums[0] ** 2
Perfect square n (4, 9, 16, 25, 36, 49) → i == n / i at the top, and the guard prevents double counting
Prime n → only indices 1 and n qualify
Index 1 always divides n, and index n always does, so the answer is never 0
Maximum value: n = 50 with every element 50 gives at most 6 divisors of 50 times 2500, well inside int

Code

from typing import List


class Solution:
    def sumOfSquares(self, nums: List[int]) -> int:
        """
        2778. Sum of Squares of Special Elements
        Time: O(sqrt(n))
        Space: O(1)
        """
        n = len(nums)
        total = 0

        # the special positions are exactly the divisors of n, so enumerate
        # those in pairs rather than scanning every index
        i = 1
        while i * i <= n:
            if n % i == 0:
                total += nums[i - 1] ** 2  # statement is 1-indexed
                j = n // i
                if j != i:
                    total += nums[j - 1] ** 2
            i += 1

        return total

/**
 * 2778. Sum of Squares of Special Elements
 * Time: O(sqrt(n))
 * Space: O(1)
 */
class Solution {
    public int sumOfSquares(int[] nums) {
        int n = nums.length;
        int total = 0;

        // the special positions are exactly the divisors of n, so enumerate
        // those in pairs rather than scanning every index
        for (int i = 1; i * i <= n; i++) {
            if (n % i != 0) {
                continue;
            }
            total += nums[i - 1] * nums[i - 1];  // statement is 1-indexed
            int j = n / i;
            if (j != i) {
                total += nums[j - 1] * nums[j - 1];
            }
        }

        return total;
    }
}