Subarrays Distinct Element Sum of Squares II solution and explanation
2 more videos on this problem
Beautiful Segment tree lazy idea | Subarrays Distinct Element Sum of Squares IIVivek Gupta · 24:19 · 4K views
2916. Subarrays Distinct Element Sum of Squares II | 2913 | Leetcode Biweekly 116codingMohan · 37:51 · 3K views
2916. Subarrays Distinct Element Sum of Squares II (Leetcode Hard)Programming Live with Larry · 2:25:46 · 475 views
The problem
You are given a 0-indexed integer array nums.
The distinct count of a subarray of nums is defined as:
- Let
nums[i..j]be a subarray ofnumsconsisting of all the indices fromitojsuch that0 <= i <= j < nums.length. Then the number of distinct values innums[i..j]is called the distinct count ofnums[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] <= 105
Approach overview
Problem 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.
Approach 1: Brute Force Approach
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.
def sum_of_squares(nums):
MOD = 10**9 + 7
total_sum = 0
n = len(nums)
for i in range(n):
distinct_count = 0
seen = set()
for j in range(i, n):
if nums[j] not in seen:
distinct_count += 1
seen.add(nums[j])
total_sum += distinct_count ** 2
total_sum %= MOD
return total_sum#include <vector>
#include <unordered_set>
using namespace std;
int sumOfSquares(vector<int>& nums) {
const int MOD = 1000000007;
int total_sum = 0, n = nums.size();
for (int i = 0; i < n; ++i) {
unordered_set<int> seen;
int distinct_count = 0;
for (int j = i; j < n; ++j) {
if (seen.find(nums[j]) == seen.end()) {
seen.insert(nums[j]);
++distinct_count;
}
total_sum = (total_sum + distinct_count * distinct_count) % MOD;
}
}
return total_sum;
}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.
Approach 2: Sliding Window Technique
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.
def sum_of_squares_eff(nums):
MOD = 10**9 + 7
n = len(nums)
left = 0
total_sum = 0
count = [0] * 100001
unique_count = 0
for right in range(n):
count[nums[right]] += 1
if count[nums[right]] == 1:
unique_count += 1
while left <= right and count[nums[right]] > 1:
count[nums[left]] -= 1
if count[nums[left]] == 0:
unique_count -= 1
left += 1
total_sum += unique_count * unique_count
total_sum %= MOD
return total_sum#include <vector>
#include <unordered_map>
using namespace std;
int sumOfSquaresSliding(vector<int>& nums) {
const int MOD = 1000000007;
int n = nums.size();
unordered_map<int, int> count;
int total_sum = 0, unique_count = 0;
int left = 0;
for (int right = 0; right < n; ++right) {
count[nums[right]]++;
if (count[nums[right]] == 1)
unique_count++;
while (left <= right && count[nums[right]] > 1) {
count[nums[left]]--;
if (count[nums[left]] == 0)
unique_count--;
left++;
}
total_sum = (total_sum + unique_count * unique_count) % MOD;
}
return total_sum;
}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.
Complexity comparison
| Approach | Complexity |
|---|---|
| Brute Force Approach | Time Complexity: O(n3) - Since each possible subarray is iterated over. |
| Sliding Window Technique | Time Complexity: O(n) |
Complexity analysis
| 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 |
Frequently asked questions
Is Subarrays Distinct Element Sum of Squares II easy or hard?
Subarrays Distinct Element Sum of Squares II is classified as a Hard problem. It requires understanding subarray contribution patterns, tracking last occurrences, and implementing advanced data structures such as Fenwick Trees or Segment Trees to achieve the optimal O(n log n) solution.
Subarrays Distinct Element Sum of Squares II Python/Java solution
Python and Java implementations typically use a Fenwick Tree or Segment Tree combined with a dictionary or array to track last occurrences. Each iteration updates contribution ranges and queries cumulative values to maintain the total sum of squared distinct counts.
How to solve Subarrays Distinct Element Sum of Squares II in O(n log n)?
Process the array from left to right while tracking the previous index of each value. When a value appears again, adjust the contribution range where it affects distinct counts. Use a Fenwick Tree or Segment Tree to apply range updates and compute prefix sums of squared distinct counts efficiently.
What is the best approach for Subarrays Distinct Element Sum of Squares II?
The most efficient solution uses a contribution-based technique with a Binary Indexed Tree (Fenwick Tree) or Segment Tree. While scanning the array, track the last occurrence of each element and update the range of subarrays whose distinct counts change. This reduces repeated work and achieves O(n log n) time complexity.
Is Subarrays Distinct Element Sum of Squares II asked at Google/Amazon/Meta?
Hard array and data structure problems involving segment trees or Fenwick trees frequently appear in interviews at companies like Google, Amazon, and Meta. This problem tests advanced subarray contribution techniques and efficient range updates, which are common interview patterns.
What data structure is used in Subarrays Distinct Element Sum of Squares II?
Efficient solutions rely on Binary Indexed Trees (Fenwick Trees) or Segment Trees to handle range updates and prefix queries. A hash map or array is also used to store the last occurrence index of each element while scanning the array.
What is the time complexity of Subarrays Distinct Element Sum of Squares II?
The brute force method runs in O(n^2) time because it enumerates every subarray and maintains a set of distinct elements. The optimized solution processes the array once and performs logarithmic range updates and queries using Fenwick or segment trees, resulting in O(n log n) time and O(n) space.