# Subarrays Distinct Element Sum of Squares II
**Difficulty:** HARD
[External](https://leetcode.com/problems/subarrays-distinct-element-sum-of-squares-ii)
Canonical: https://scaleengineer.com/dsa/problems/subarrays-distinct-element-sum-of-squares-ii
**Patterns:** [Dynamic Programming](https://scaleengineer.com/dsa/patterns/dynamic-programming)
**Data structures:** Array, Binary Indexed Tree, Segment Tree
---
## 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 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 <= 105`
* `1 <= nums[i] <= 105`

# Approaches
## Brute Force Enumeration
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.
**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.
**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.
### Explanation
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`.

```java
import java.util.HashSet;

class Solution {
    public int sumCounts(int[] nums) {
        int n = nums.length;
        long totalSum = 0;
        int MOD = 1_000_000_007;

        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                HashSet<Integer> distinctElements = new HashSet<>();
                for (int k = i; k <= j; k++) {
                    distinctElements.add(nums[k]);
                }
                long distinctCount = distinctElements.size();
                totalSum = (totalSum + (distinctCount * distinctCount)) % MOD;
            }
        }
        return (int) totalSum;
    }
}
```
### 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`.

## Optimized Brute Force
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.
**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`.
**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).
### Explanation
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.

```java
import java.util.HashSet;

class Solution {
    public int sumCounts(int[] nums) {
        int n = nums.length;
        long totalSum = 0;
        int MOD = 1_000_000_007;

        for (int i = 0; i < n; i++) {
            HashSet<Integer> distinctElements = new HashSet<>();
            for (int j = i; j < n; j++) {
                distinctElements.add(nums[j]);
                long distinctCount = distinctElements.size();
                totalSum = (totalSum + (distinctCount * distinctCount)) % MOD;
            }
        }
        return (int) totalSum;
    }
}
```
### 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`.

## Dynamic Programming with Segment Tree
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.
**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.
**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.
### Explanation
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:
1. `sum_sq`: The sum of squares of values in its range.
2. `sum_lin`: The linear sum of values in its range.
3. `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]`.

```java
import java.util.HashMap;

class Solution {
    private static final int MOD = 1_000_000_007;
    private long[] sumSq, sumLin, lazy;
    private int n;

    public int sumCounts(int[] nums) {
        this.n = nums.length;
        sumSq = new long[4 * n];
        sumLin = new long[4 * n];
        lazy = new long[4 * n];

        long totalSum = 0;
        long currentSumOfSquares = 0;
        HashMap<Integer, Integer> lastSeen = new HashMap<>();

        for (int j = 0; j < n; j++) {
            int p = lastSeen.getOrDefault(nums[j], -1);
            update(1, 0, n - 1, p + 1, j, 1);
            currentSumOfSquares = sumSq[1];
            totalSum = (totalSum + currentSumOfSquares) % MOD;
            lastSeen.put(nums[j], j);
        }

        return (int) totalSum;
    }

    private void pushDown(int node, int start, int end) {
        if (lazy[node] == 0) return;

        long c = lazy[node];
        int mid = start + (end - start) / 2;
        int leftChild = 2 * node;
        int rightChild = 2 * node + 1;

        // Update left child
        long leftLen = mid - start + 1;
        sumSq[leftChild] = (sumSq[leftChild] + 2 * c * sumLin[leftChild] + (c * c % MOD) * leftLen) % MOD;
        sumLin[leftChild] = (sumLin[leftChild] + c * leftLen) % MOD;
        lazy[leftChild] = (lazy[leftChild] + c) % MOD;

        // Update right child
        long rightLen = end - mid;
        sumSq[rightChild] = (sumSq[rightChild] + 2 * c * sumLin[rightChild] + (c * c % MOD) * rightLen) % MOD;
        sumLin[rightChild] = (sumLin[rightChild] + c * rightLen) % MOD;
        lazy[rightChild] = (lazy[rightChild] + c) % MOD;

        lazy[node] = 0;
    }

    private void update(int node, int start, int end, int l, int r, int val) {
        if (start > end || start > r || end < l) {
            return;
        }

        if (l <= start && end <= r) {
            long c = val;
            long len = end - start + 1;
            sumSq[node] = (sumSq[node] + 2 * c * sumLin[node] + (c * c % MOD) * len) % MOD;
            sumLin[node] = (sumLin[node] + c * len) % MOD;
            if (start != end) {
                lazy[node] = (lazy[node] + c) % MOD;
            }
            return;
        }

        pushDown(node, start, end);

        int mid = start + (end - start) / 2;
        update(2 * node, start, mid, l, r, val);
        update(2 * node + 1, mid + 1, end, l, r, val);

        sumSq[node] = (sumSq[2 * node] + sumSq[2 * node + 1]) % MOD;
        sumLin[node] = (sumLin[2 * node] + sumLin[2 * node + 1]) % MOD;
    }
}
```
### 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`.