You are given a 0-indexed integer array nums.
The distinct count of a subarray of nums is defined as:
nums[i..j] be a subarray of nums consisting of all the indices from i to j such that 0 <= i <= j < nums.length. Then the number of distinct values in nums[i..j] is called the distinct count of nums[i..j].Return the sum of the squares of distinct counts of all subarrays of nums.
Since the answer may be very large, return it modulo 109 + 7.
A subarray is a contiguous non-empty sequence of elements within an array.
Example 1:
Input: nums = [1,2,1] Output: 15 Explanation: Six possible subarrays are: [1]: 1 distinct value [2]: 1 distinct value [1]: 1 distinct value [1,2]: 2 distinct values [2,1]: 2 distinct values [1,2,1]: 2 distinct values The sum of the squares of the distinct counts in all subarrays is equal to 12 + 12 + 12 + 22 + 22 + 22 = 15.
Example 2:
Input: nums = [2,2] Output: 3 Explanation: Three possible subarrays are: [2]: 1 distinct value [2]: 1 distinct value [2,2]: 1 distinct value The sum of the squares of the distinct counts in all subarrays is equal to 12 + 12 + 12 = 3.
Constraints:
1 <= nums.length <= 1051 <= nums[i] <= 105Problem Overview: Given an array nums, evaluate every subarray and count how many distinct elements it contains. Square that count and sum the result across all subarrays. The challenge is avoiding the naive O(n²) enumeration of distinct counts.
Approach 1: Brute Force with Set Tracking (O(n²) time, O(n) space)
Start a subarray at index i and expand the right boundary one step at a time. Maintain a set or frequency map to track distinct elements as you iterate j = i..n-1. After each extension, compute the current number of unique elements and add distinct² to the answer. This approach directly simulates the definition of the problem and is easy to reason about, but repeated scans of overlapping subarrays make it quadratic.
Approach 2: Sliding Window + Contribution Tracking (O(n log n) time, O(n) space)
The optimized solution tracks how each new element affects the number of distinct elements across multiple subarrays simultaneously. Maintain the last occurrence of each value and compute how many subarrays ending at the current index gain a new distinct element. Instead of recomputing counts for every subarray, accumulate contributions using a data structure such as a Fenwick Tree or Segment Tree. These structures efficiently update ranges and query aggregated values while processing the array left to right. Each element updates the contribution range where it becomes a new distinct value, and the squared distinct counts are aggregated incrementally.
This approach relies on range updates and prefix queries, which makes data structures from Binary Indexed Tree and Segment Tree especially useful. The state transitions resemble a dynamic contribution model often seen in advanced Dynamic Programming problems on subarrays.
Recommended for interviews: Start by explaining the brute force approach to show you understand the definition of the problem. Then move to the optimized contribution-based solution using Fenwick or segment trees. Interviewers typically expect the O(n log n) approach because it demonstrates strong control over range updates, prefix queries, and subarray contribution analysis.
This method involves generating all possible subarrays of the given array and calculating the number of distinct elements in each subarray. The distinct count for each subarray is then squared and summed up to get the final result.
The Python solution initializes a total_sum to zero. It then iterates over all possible start points of subarrays. For each start, it calculates the distinct count while moving to different end points, storing encountered numbers in a set. The square of this distinct count is added to the total_sum, which is returned after taking modulo 109 + 7.
Time Complexity: O(n3) - Since each possible subarray is iterated over.
Space Complexity: O(n) - For storing distinct elements using a set.
Using the Sliding Window Technique can improve efficiency by reducing repeated calculations of distinct counts. By maintaining a sliding window over the array, we can update the distinct count as the window slides, maintaining a dynamic result set of numbers.
This Python solution utilizes a sliding window technique with two pointers, left and right, representing the current window. A frequency array 'count' keeps track of element occurrences, and unique_count stores the number of distinct elements in the current window. After each iteration, the distinct count squared is added to total_sum.
Time Complexity: O(n)
Space Complexity: O(1) when considering the auxiliary space for counting and a fixed-size integer array.
| Approach | Complexity |
|---|---|
| Brute Force Approach | Time Complexity: O(n3) - Since each possible subarray is iterated over. |
| Sliding Window Technique | Time Complexity: O(n) |
| Approach | Time | Space | When to Use |
|---|---|---|---|
| Brute Force with Set | O(n²) | O(n) | Small arrays or when demonstrating the base idea of counting distinct elements per subarray |
| Sliding Window + Fenwick/Segment Tree | O(n log n) | O(n) | Large inputs where efficient range updates and prefix aggregation are required |
Beautiful Segment tree lazy idea | Subarrays Distinct Element Sum of Squares II • Vivek Gupta • 4,403 views views
Watch 3 more video solutions →Practice Subarrays Distinct Element Sum of Squares II with our built-in code editor and test cases.
Practice on FleetCode