Problem Description

You have an integer array nums that is 0-indexed.

For any subarray of nums, you need to find its distinct count, which is the number of unique/distinct values present in that subarray. A subarray nums[i..j] includes all elements from index i to index j (inclusive), where 0 <= i <= j < nums.length.

Your task is to:

1. Find all possible subarrays of nums
2. Calculate the distinct count for each subarray
3. Square each distinct count
4. Return the sum of all these squared values

For example, if a subarray has 3 distinct values, its contribution to the final answer would be 3 * 3 = 9.

Key Points:

• A subarray must be contiguous (elements next to each other in the original array)
• A subarray cannot be empty
• You need to consider all possible subarrays
• The distinct count is the number of unique elements in a subarray (duplicates count as one)

The solution approach iterates through all possible starting positions i, and for each starting position, extends the subarray by incrementing the ending position j. As each new element nums[j] is added to the current subarray, it's added to a set s which automatically handles uniqueness. The size of this set represents the distinct count, and len(s) * len(s) gives the square of the distinct count, which is added to the running total.

Intuition

To solve this problem, we need to examine every possible subarray and calculate its distinct count. The most straightforward way to think about this is to consider how we can systematically generate all subarrays.

A subarray is defined by two things: where it starts and where it ends. If we fix a starting position i, we can extend the subarray one element at a time by moving the ending position j from i to the end of the array. This gives us all subarrays that start at position i.

The key insight is that when we extend a subarray by one element (moving from nums[i..j] to nums[i..j+1]), we're just adding one new element nums[j+1] to our existing subarray. This means we don't need to recalculate the distinct count from scratch each time - we can maintain it incrementally.

A set is the perfect data structure for tracking distinct elements because:

• It automatically handles duplicates (adding the same element twice doesn't change the set)
• We can query its size in constant time to get the distinct count
• Adding an element is efficient

So the approach becomes:

1. For each starting position i, create an empty set
2. Extend the subarray one element at a time (incrementing j)
3. Add each new element to the set
4. The size of the set at any point gives us the distinct count for the current subarray
5. Square this count and add it to our answer

This way, we process each subarray exactly once, and for each subarray, we efficiently calculate its contribution (distinct_count²) to the final sum.

Solution Approach

The solution uses a nested loop approach with a set data structure to efficiently track distinct elements.

Algorithm Steps:

1. Initialize variables: Set ans = 0 to store the final sum and get the length of the array n = len(nums).

2. Outer loop - Fix the starting position: Iterate through each possible starting index i from 0 to n-1.

3. Initialize a set for each starting position: For each i, create an empty set s = set(). This will track the distinct elements in subarrays starting at index i.

4. Inner loop - Extend the subarray: For each starting position i, iterate j from i to n-1. This represents extending the subarray one element at a time.

5. Update the set and calculate contribution:

   • Add the current element nums[j] to the set: s.add(nums[j])
   • The set automatically handles duplicates, so len(s) gives us the current distinct count
   • Calculate the square of the distinct count: len(s) * len(s)
   • Add this squared value to our answer: ans += len(s) * len(s)

6. Return the result: After processing all subarrays, return ans.

Why this works:

• When j = i, we have a subarray with just one element nums[i]
• As j increases, we're considering subarrays nums[i..i+1], nums[i..i+2], and so on
• The set maintains the distinct elements seen so far in the current subarray
• By adding len(s) * len(s) at each step, we're accounting for each subarray exactly once

Time Complexity: O(n²) where n is the length of the array, as we have two nested loops.

Space Complexity: O(n) in the worst case when all elements in a subarray are distinct.

Example Walkthrough

Let's walk through the solution with nums = [1, 2, 1].

We need to find all subarrays, calculate their distinct counts, square them, and sum the results.

Step-by-step execution:

When i = 0 (starting at index 0):

• Initialize empty set s = {}
• j = 0: Add nums[0] = 1 to set → s = {1}
  • Subarray [1] has distinct count = 1
  • Contribution: 1² = 1
  • Running total: ans = 1
• j = 1: Add nums[1] = 2 to set → s = {1, 2}
  • Subarray [1, 2] has distinct count = 2
  • Contribution: 2² = 4
  • Running total: ans = 1 + 4 = 5
• j = 2: Add nums[2] = 1 to set → s = {1, 2} (1 already exists)
  • Subarray [1, 2, 1] has distinct count = 2
  • Contribution: 2² = 4
  • Running total: ans = 5 + 4 = 9

When i = 1 (starting at index 1):

• Initialize empty set s = {}
• j = 1: Add nums[1] = 2 to set → s = {2}
  • Subarray [2] has distinct count = 1
  • Contribution: 1² = 1
  • Running total: ans = 9 + 1 = 10
• j = 2: Add nums[2] = 1 to set → s = {2, 1}
  • Subarray [2, 1] has distinct count = 2
  • Contribution: 2² = 4
  • Running total: ans = 10 + 4 = 14

When i = 2 (starting at index 2):

• Initialize empty set s = {}
• j = 2: Add nums[2] = 1 to set → s = {1}
  • Subarray [1] has distinct count = 1
  • Contribution: 1² = 1
  • Running total: ans = 14 + 1 = 15

Final answer: 15

This demonstrates how the set efficiently tracks distinct elements as we extend each subarray, and how we calculate the squared distinct count for each subarray's contribution to the final sum.

Solution Implementation

Time and Space Complexity

The time complexity is O(n^2), where n is the length of the array nums. This is because we have two nested loops: the outer loop runs n times (iterating through each starting position i), and for each iteration of the outer loop, the inner loop runs from position i to n-1, which on average contributes to n/2 iterations. The total number of operations is approximately n * n/2, which simplifies to O(n^2).

The space complexity is O(n). This is determined by the set s which, in the worst case, can store all unique elements from a subarray. Since the maximum number of unique elements in any subarray cannot exceed the total number of elements in the array, the set can grow up to size n. The other variables (ans, n, i, j) use constant space O(1), so the overall space complexity is dominated by the set, resulting in O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Confusing Subarrays with Subsequences

A frequent mistake is treating this problem as finding all possible subsequences instead of subarrays. Remember that subarrays must be contiguous (consecutive elements), while subsequences can skip elements.

Incorrect approach:

# Wrong: This would generate subsequences, not subarrays
def wrong_approach(nums):
    # Trying to use combinations or bit manipulation to get all subsets
    from itertools import combinations
    total = 0
    for length in range(1, len(nums) + 1):
        for subset in combinations(nums, length):
            distinct = len(set(subset))
            total += distinct * distinct
    return total

Correct approach: Use nested loops where the inner loop extends from the current starting position, ensuring contiguity.

2. Recreating the Set for Each Subarray

Some might think they need to recreate the set from scratch for every single subarray, leading to inefficient code:

Inefficient approach:

def inefficient_approach(nums):
    total_sum = 0
    n = len(nums)
    for i in range(n):
        for j in range(i, n):
            # Creating a new set for each subarray nums[i:j+1]
            distinct_elements = set(nums[i:j+1])  # O(j-i+1) operation
            total_sum += len(distinct_elements) ** 2
    return total_sum

This approach has O(n³) time complexity because slicing and creating a set for each subarray takes O(n) time in the worst case.

Optimized approach: Build the set incrementally as shown in the original solution, which maintains O(n²) complexity.

3. Integer Overflow in Other Languages

While Python handles large integers automatically, in languages like Java or C++, squaring the distinct count and summing could cause integer overflow.

Solution for other languages:

• Use long or long long data types
• Consider using modulo arithmetic if the problem specifies it

// Java example
long totalSum = 0;
// ... 
totalSum += (long) distinctCount * distinctCount;

4. Off-by-One Errors in Loop Boundaries

It's easy to make mistakes with the loop boundaries, especially forgetting that the end index should be inclusive:

Common mistake:

# Wrong: Missing the last element in some subarrays
for end_index in range(start_index, array_length - 1):  # Should be array_length
    distinct_elements.add(nums[end_index])

Correct: Ensure the inner loop goes up to array_length - 1 (inclusive) by using range(start_index, array_length).