LeetCode 2913 - Subarrays Distinct Element Sum of Squares I

LeetCode Problem 2913

Difficulty: 🟢 Easy
Topics: Array, Hash Table

Solution

Problem Understanding

The problem asks us to examine every possible non-empty subarray of the given array nums. For each subarray, we compute how many distinct values appear inside it. After finding this distinct count, we square it and add it to the final answer.

A subarray is contiguous, which means the elements must appear next to each other in the original array. For an array of length n, there are n * (n + 1) / 2 total subarrays.

For example, if the subarray is [1,2,1], the distinct values are {1,2}, so the distinct count is 2. The contribution of this subarray to the answer is 2² = 4.

The input consists of:

• An integer array nums
• nums[i] represents the value at index i

The output is:

• A single integer representing the sum of the squares of distinct counts for every possible subarray

The constraints are small:

• 1 <= nums.length <= 100
• 1 <= nums[i] <= 100

These limits are important because they tell us that an O(n²) or even O(n³) solution is acceptable. Since the array length is at most 100, brute-force enumeration of subarrays is feasible.

There are several edge cases worth considering:

• Arrays with all identical values, where every subarray has distinct count 1
• Arrays with all unique values, where distinct counts grow as the subarray expands
• Arrays of length 1, where the answer is always 1
• Repeated patterns such as [1,2,1,2], where distinct counts increase and then stabilize

The problem guarantees that the array is non-empty, so we never need to handle an empty input.

Approaches

Brute Force Approach

The most direct solution is to generate every possible subarray and explicitly compute how many distinct elements it contains.

For every starting index i, we can try every ending index j. For each subarray nums[i:j+1], we create a set containing its elements. The size of the set gives the distinct count. We square that count and add it to the answer.

This approach is correct because it evaluates every subarray independently and computes the exact number of distinct elements.

However, if we rebuild the set from scratch for every subarray, the complexity becomes expensive. Constructing the set for one subarray may take O(n) time, and there are O(n²) subarrays, leading to O(n³) total time complexity.

Although n <= 100 makes even this acceptable, we can do better.

Improved Incremental Hash Set Approach

The key observation is that when we extend a subarray by one element, we do not need to recompute the distinct elements from scratch.

Suppose we fix a starting index i. As we move the ending index j from left to right, we can maintain a hash set containing all distinct values currently in the subarray.

Each time we add nums[j]:

• Insert it into the set
• The set automatically keeps only unique values
• The size of the set is the current distinct count

This avoids rebuilding the set repeatedly. Each subarray extension only performs one insertion into the hash set.

This reduces the time complexity from O(n³) to O(n²).

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n³)	O(n)	Rebuilds the distinct set for every subarray
Optimal	O(n²)	O(n)	Maintains a running hash set while extending subarrays

Algorithm Walkthrough

1. Initialize a variable answer to store the total sum.
2. Iterate through every possible starting index start from 0 to n - 1.
3. For each starting index, create an empty hash set called distinct_values.
4. Extend the subarray one element at a time using an ending index end.
5. Insert nums[end] into the hash set.

This works because sets automatically remove duplicates. If the value already exists, the set remains unchanged. 6. Compute the current distinct count using len(distinct_values). 7. Square the distinct count and add it to answer. 8. Continue extending the subarray until all ending positions have been processed. 9. After processing all starting indices, return answer.

Why it works

For every starting index, the algorithm incrementally builds all possible subarrays beginning at that position. The hash set always contains exactly the distinct elements in the current subarray. Since every subarray is visited exactly once and its distinct count is computed correctly, the final accumulated sum is correct.

Python Solution

from typing import List

class Solution:
    def sumCounts(self, nums: List[int]) -> int:
        n = len(nums)
        answer = 0

        for start in range(n):
            distinct_values = set()

            for end in range(start, n):
                distinct_values.add(nums[end])

                distinct_count = len(distinct_values)
                answer += distinct_count * distinct_count

        return answer

The implementation follows the algorithm directly.

The outer loop selects the starting index of the subarray. For every new starting position, we reset the hash set because we are beginning a completely different collection of subarrays.

The inner loop extends the subarray one element at a time. Each new element is inserted into the set, which maintains only unique values.

After updating the set, we compute the number of distinct elements using len(distinct_values). We square this value and add it to the running total.

Because the set is updated incrementally instead of rebuilt repeatedly, the solution remains efficient.

Go Solution

func sumCounts(nums []int) int {
	n := len(nums)
	answer := 0

	for start := 0; start < n; start++ {
		distinctValues := make(map[int]bool)

		for end := start; end < n; end++ {
			distinctValues[nums[end]] = true

			distinctCount := len(distinctValues)
			answer += distinctCount * distinctCount
		}
	}

	return answer
}

The Go implementation mirrors the Python logic closely.

Since Go does not have a built-in set type, we use a map[int]bool to simulate a hash set. The keys represent the distinct values currently present in the subarray.

The built-in len() function on the map gives the number of distinct elements.

Integer overflow is not a concern here because the constraints are very small. Even the maximum possible answer easily fits inside a standard Go int.

Worked Examples

Example 1

Input:

nums = [1,2,1]

We process all subarrays.

Start = 0

End	Subarray	Distinct Set	Distinct Count	Contribution	Running Total
0	[1]	{1}	1	1	1
1	[1,2]	{1,2}	2	4	5
2	[1,2,1]	{1,2}	2	4	9

Start = 1

End	Subarray	Distinct Set	Distinct Count	Contribution	Running Total
1	[2]	{2}	1	1	10
2	[2,1]	{1,2}	2	4	14

Start = 2

End	Subarray	Distinct Set	Distinct Count	Contribution	Running Total
2	[1]	{1}	1	1	15

Final answer:

15

Example 2

Input:

nums = [1,1]

Start = 0

End	Subarray	Distinct Set	Distinct Count	Contribution	Running Total
0	[1]	{1}	1	1	1
1	[1,1]	{1}	1	1	2

Start = 1

End	Subarray	Distinct Set	Distinct Count	Contribution	Running Total
1	[1]	{1}	1	1	3

Final answer:

3

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n²)	We examine every subarray once
Space	O(n)	The hash set may store up to n distinct elements

There are O(n²) subarrays in total. For each subarray extension, we perform constant-time hash set operations. Therefore, the total runtime is O(n²).

The extra space comes from the hash set used to track distinct elements in the current subarray. In the worst case, all elements are unique, so the set may contain n elements.

Test Cases

from typing import List

class Solution:
    def sumCounts(self, nums: List[int]) -> int:
        n = len(nums)
        answer = 0

        for start in range(n):
            distinct_values = set()

            for end in range(start, n):
                distinct_values.add(nums[end])
                answer += len(distinct_values) ** 2

        return answer

sol = Solution()

assert sol.sumCounts([1, 2, 1]) == 15  # example with repeated value
assert sol.sumCounts([1, 1]) == 3  # all subarrays have one distinct value
assert sol.sumCounts([5]) == 1  # single element array
assert sol.sumCounts([1, 2, 3]) == 20  # all elements unique
assert sol.sumCounts([2, 2, 2, 2]) == 10  # all identical values
assert sol.sumCounts([1, 2, 1, 2]) == 28  # repeating alternating pattern
assert sol.sumCounts([1, 2]) == 6  # small two-element unique array
assert sol.sumCounts([3, 3, 4]) == 9  # duplicate followed by new value

Test Summary

Test	Why
[1,2,1]	Validates repeated values and changing distinct counts
[1,1]	Ensures duplicates do not increase distinct count
[5]	Smallest valid input
[1,2,3]	Tests all unique elements
[2,2,2,2]	Tests constant distinct count
[1,2,1,2]	Tests alternating repeated values
[1,2]	Simple unique two-element case
[3,3,4]	Tests transition from duplicate to new distinct value

Edge Cases

One important edge case is an array containing only one element. In this scenario, there is exactly one subarray, and its distinct count is always 1. The implementation handles this naturally because the nested loops execute once and compute 1².

Another important case is when all elements are identical, such as [2,2,2,2]. Every subarray still has only one distinct value, regardless of its length. A buggy implementation might incorrectly count duplicates multiple times. Using a hash set prevents this issue because duplicate insertions do not change the set size.

A third important edge case is when all elements are unique, such as [1,2,3,4]. In this case, the number of distinct elements increases every time the subarray expands. The implementation handles this correctly because each new value added to the set increases its size by one.

A more subtle case is repeating patterns like [1,2,1,2]. The distinct count sometimes increases and sometimes remains unchanged depending on whether the newly added element already exists in the current subarray. The hash set accurately captures this behavior because inserting an existing value leaves the set unchanged.