Subarrays Distinct Element Sum of Squares II

Hard

#2597Time: O(n³) - There are three nested loops. The outer two loops iterate through all O(n²) subarrays, and for each subarray of average length O(n), we iterate through its elements. This results in a cubic time complexity.Space: O(n) - In the worst-case scenario, a subarray can contain `n` distinct elements, requiring the `HashSet` to store up to `n` items.

Patterns

Dynamic Programming

Data structures

Array Binary Indexed Tree Segment Tree

Prompt

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 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 10⁹ + 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 <= 10⁵
1 <= nums[i] <= 10⁵

Approaches

3 approaches with complexity analysis and trade-offs.

This is the most straightforward and naive approach. It involves generating every possible subarray, and for each subarray, counting the number of distinct elements. The square of this count is then added to a running total. This method is easy to conceptualize but highly inefficient.

Algorithm

Initialize total_sum to 0.
Iterate through all possible start indices i from 0 to n-1.
Iterate through all possible end indices j from i to n-1.
- This defines a subarray nums[i..j].
- Create a new HashSet to store elements of this subarray.
- Iterate from k = i to j and add nums[k] to the HashSet.
- Get the size of the HashSet, which is the distinct count d.
- Add d*d to total_sum, taking modulo 10^9 + 7.
Return total_sum.

Walkthrough

The brute-force approach uses three nested loops to solve the problem. The outer two loops, with indices i and j, are used to define the start and end of every possible subarray nums[i..j]. For each of these subarrays, a third loop iterates from i to j. Inside this innermost loop, we use a HashSet to keep track of the unique elements encountered within the current subarray. After the inner loop completes, the size of the HashSet gives us the distinct count, d. We then compute d*d and add it to our total sum. To prevent integer overflow, all additions to the sum are performed modulo 10^9 + 7.

1import java.util.HashSet;2 3class Solution {4    public int sumCounts(int[] nums) {5        int n = nums.length;6        long totalSum = 0;7        int MOD = 1_000_000_007;8 9        for (int i = 0; i < n; i++) {10            for (int j = i; j < n; j++) {11                HashSet<Integer> distinctElements = new HashSet<>();12                for (int k = i; k <= j; k++) {13                    distinctElements.add(nums[k]);14                }15                long distinctCount = distinctElements.size();16                totalSum = (totalSum + (distinctCount * distinctCount)) % MOD;17            }18        }19        return (int) totalSum;20    }21}

Complexity

Time

O(n³) - There are three nested loops. The outer two loops iterate through all O(n²) subarrays, and for each subarray of average length O(n), we iterate through its elements. This results in a cubic time complexity.

Space

O(n) - In the worst-case scenario, a subarray can contain `n` distinct elements, requiring the `HashSet` to store up to `n` items.

Trade-offs

Pros

Simple to understand and implement.
Correct for small input sizes.

Cons

Extremely inefficient due to its cubic time complexity.
Will result in a 'Time Limit Exceeded' error for the given constraints.

This approach is an optimization of the naive brute-force method. Instead of recalculating the distinct elements for each subarray from scratch, we can fix the starting point of a subarray and extend it one element at a time, updating the count of distinct elements incrementally.

Algorithm

Initialize total_sum to 0.
Iterate through all possible start indices i from 0 to n-1.
- Create an empty HashSet for subarrays starting at i.
- Iterate through all possible end indices j from i to n-1.
  - Add nums[j] to the HashSet.
  - The size of the HashSet is the distinct count d for nums[i..j].
  - Add d*d to total_sum, taking modulo 10^9 + 7.
Return total_sum.

Walkthrough

We use two nested loops. The outer loop fixes the starting index i of our subarrays. For each i, we initialize a HashSet. The inner loop then iterates from j = i to n-1, effectively considering all subarrays that start at i (nums[i..i], nums[i..i+1], etc.). As we extend the subarray by including nums[j], we add it to our HashSet. The size of the set at each step j gives the distinct count for the subarray nums[i..j]. We square this count and add it to the total sum. This avoids the third loop of the naive approach, reducing the complexity.

1import java.util.HashSet;2 3class Solution {4    public int sumCounts(int[] nums) {5        int n = nums.length;6        long totalSum = 0;7        int MOD = 1_000_000_007;8 9        for (int i = 0; i < n; i++) {10            HashSet<Integer> distinctElements = new HashSet<>();11            for (int j = i; j < n; j++) {12                distinctElements.add(nums[j]);13                long distinctCount = distinctElements.size();14                totalSum = (totalSum + (distinctCount * distinctCount)) % MOD;15            }16        }17        return (int) totalSum;18    }19}

Complexity

Time

O(n²) - Two nested loops run, and the `HashSet` add operation takes, on average, O(1) time.

Space

O(n) - The `HashSet` can store up to `n` distinct elements in the worst case for a given starting index `i`.

Trade-offs

Pros

Improves upon the O(n³) brute force.
Still relatively easy to implement.

Cons

The quadratic time complexity is still too slow for the given constraints (n up to 10^5).

This efficient approach uses dynamic programming combined with a segment tree. We process the array from left to right. At each index j, we calculate the sum of squares of distinct counts for all subarrays ending at j. A key observation is that when we move from index j-1 to j, the distinct counts of subarrays nums[i..j] change in a predictable way based on the last occurrence of nums[j]. This pattern allows for a range update on the distinct counts, which can be handled efficiently by a segment tree.

Algorithm

Initialize a segment tree of size n with all values 0. Each node will store sum_sq (sum of squares), sum_lin (linear sum), and a lazy value for range updates.
Initialize total_sum = 0 and a prev map to store the last seen index of each number.
Loop j from 0 to n-1:
- Get the previous index p of nums[j] from the prev map (or -1 if not seen).
- Perform a range increment operation on the segment tree for the index range [p+1, j], adding 1 to each value.
- The sum of squares for all subarrays ending at j, let's call it dp_j, is now the total sum of squares stored in the root of the segment tree.
- Add dp_j to total_sum (modulo 10^9 + 7).
- Update prev[nums[j]] to j.
Return total_sum.

Walkthrough

Let d(i, j) be the distinct count of nums[i..j]. At each step j, we are interested in dp[j] = sum_{i=0}^{j} d(i, j)^2. The total answer is sum_{j=0}^{n-1} dp[j].

Let p be the previous index of nums[j]. For any subarray nums[i..j], if i > p, nums[j] is a new element, so d(i, j) = d(i, j-1) + 1. If i <= p, nums[j] was already present (as nums[p]), so d(i, j) = d(i, j-1).

This means that as we process nums[j], the array of distinct counts for subarrays ending at j ([d(0,j), d(1,j), ..., d(j,j)]) can be obtained by taking the array for j-1 and performing a range increment of +1 on the suffix starting at index p+1. We also set the new value for index j to 1.

We can maintain these distinct counts in a segment tree. The segment tree will support range increments and querying the total sum of squares. Each node in the tree needs to store:

sum_sq: The sum of squares of values in its range.
sum_lin: The linear sum of values in its range.
lazy: A pending value to be added to its range.

When a range is incremented by c, the new sum of squares becomes sum(v+c)^2 = sum(v^2 + 2cv + c^2) = sum(v^2) + 2c*sum(v) + c^2*length. This shows why we need to track both sum_sq and sum_lin.

The main loop iterates from j=0 to n-1, performs a range update on the segment tree for [p+1, j], queries the root for the new total sum of squares, adds it to the answer, and updates the last seen position of nums[j].

1import java.util.HashMap;2 3class Solution {4    private static final int MOD = 1_000_000_007;5    private long[] sumSq, sumLin, lazy;6    private int n;7 8    public int sumCounts(int[] nums) {9        this.n = nums.length;10        sumSq = new long[4 * n];11        sumLin = new long[4 * n];12        lazy = new long[4 * n];13 14        long totalSum = 0;15        long currentSumOfSquares = 0;16        HashMap<Integer, Integer> lastSeen = new HashMap<>();17 18        for (int j = 0; j < n; j++) {19            int p = lastSeen.getOrDefault(nums[j], -1);20            update(1, 0, n - 1, p + 1, j, 1);21            currentSumOfSquares = sumSq[1];22            totalSum = (totalSum + currentSumOfSquares) % MOD;23            lastSeen.put(nums[j], j);24        }25 26        return (int) totalSum;27    }28 29    private void pushDown(int node, int start, int end) {30        if (lazy[node] == 0) return;31 32        long c = lazy[node];33        int mid = start + (end - start) / 2;34        int leftChild = 2 * node;35        int rightChild = 2 * node + 1;36 37        // Update left child38        long leftLen = mid - start + 1;39        sumSq[leftChild] = (sumSq[leftChild] + 2 * c * sumLin[leftChild] + (c * c % MOD) * leftLen) % MOD;40        sumLin[leftChild] = (sumLin[leftChild] + c * leftLen) % MOD;41        lazy[leftChild] = (lazy[leftChild] + c) % MOD;42 43        // Update right child44        long rightLen = end - mid;45        sumSq[rightChild] = (sumSq[rightChild] + 2 * c * sumLin[rightChild] + (c * c % MOD) * rightLen) % MOD;46        sumLin[rightChild] = (sumLin[rightChild] + c * rightLen) % MOD;47        lazy[rightChild] = (lazy[rightChild] + c) % MOD;48 49        lazy[node] = 0;50    }51 52    private void update(int node, int start, int end, int l, int r, int val) {53        if (start > end || start > r || end < l) {54            return;55        }56 57        if (l <= start && end <= r) {58            long c = val;59            long len = end - start + 1;60            sumSq[node] = (sumSq[node] + 2 * c * sumLin[node] + (c * c % MOD) * len) % MOD;61            sumLin[node] = (sumLin[node] + c * len) % MOD;62            if (start != end) {63                lazy[node] = (lazy[node] + c) % MOD;64            }65            return;66        }67 68        pushDown(node, start, end);69 70        int mid = start + (end - start) / 2;71        update(2 * node, start, mid, l, r, val);72        update(2 * node + 1, mid + 1, end, l, r, val);73 74        sumSq[node] = (sumSq[2 * node] + sumSq[2 * node + 1]) % MOD;75        sumLin[node] = (sumLin[2 * node] + sumLin[2 * node + 1]) % MOD;76    }77}

Complexity

Time

O(n log n) - The main loop runs `n` times. Inside the loop, the segment tree update operation takes O(log n) time.

Space

O(n) - The segment tree requires O(n) space. The `HashMap` for `lastSeen` can store up to O(n) unique elements in the worst case.

Trade-offs

Pros

Highly efficient with O(n log n) time complexity.
Capable of handling large inputs as per the problem constraints.

Cons

Significantly more complex to understand and implement correctly.
Requires a solid understanding of advanced data structures like segment trees with lazy propagation.

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem