2916. Subarrays Distinct Element Sum of Squares II

Problem Link

Description

---

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

Solution

---

Python3

MOD = 1000000007
 
class LazySegmentTree:
    def __init__(self, N):
        self.n = N
        self.tree = [0] * (4 * self.n)
        self.lazy = [0] * (4 * self.n)
 
    def query(self, v, tl, tr, l, r):
        self.push(v, tl, tr)
 
        if tr < l or tl > r:
            return 0
 
        if tl >= l and tr <= r:
            res = self.tree[v] % MOD
            self.lazy[v] += 1
            self.push(v, tl, tr)
            return res
        else:
            tm = tl + (tr - tl) // 2
 
            res = (self.query(2 * v, tl, tm, l, min(tm, r)) + self.query(2 * v + 1, tm + 1, tr, max(tm + 1, l), r)) % MOD
            self.tree[v] = (self.tree[2 * v] + self.tree[2 * v + 1]) % MOD
 
            return res
 
    def push(self, v, tl, tr):
        if self.lazy[v] != 0:
            self.tree[v] += (tr - tl + 1) * self.lazy[v]
 
            if tl < tr:
                self.lazy[2 * v] += self.lazy[v]
                self.lazy[2 * v + 1] += self.lazy[v]
 
            self.lazy[v] = 0
 
class Solution:
    def sumCounts(self, nums):
        N = len(nums)
        st = LazySegmentTree(N)
        last_seen = defaultdict(lambda : -1)
        curr = 0
        res = 0
 
        for i in range(N):
            x = nums[i]
            last = last_seen[x] + 1
            queryLength = i - last + 1
            curr = (curr + queryLength + 2 * st.query(1, 0, N - 1, last, i)) % MOD
            res = (res + curr) % MOD
            last_seen[x] = i
 
        return res
 

C++

const int MOD = 1000000007;
 
class LazySegmentTree {
public:
    LazySegmentTree(int N) {
        n = N;
        tree.assign(4 * n, 0);
        lazy.assign(4 * n, 0);
    }
 
    int query(int v, int tl, int tr, int l, int r) {
        lazy_update(v, tl, tr);
 
        if (tr < l || tl > r) return 0;
        
        if (tl >= l && tr <= r) {
            return tree[v] % MOD;
        } else {
            int tm = tl + (tr - tl) / 2;
 
            return (query(2 * v, tl, tm, l, min(tm, r)) + query(2 * v + 1, tm + 1, tr, max(tm + 1, l), r)) % MOD;
        }
    }
 
    void lazy_update(int v, int tl, int tr) {
        if (lazy[v] != 0) {
            tree[v] += (tr - tl + 1) * lazy[v];
            tree[v] %= MOD;
 
            if (tl < tr) {
                lazy[2 * v] += lazy[v];
                lazy[2 * v] %= MOD;
                lazy[2 * v + 1] += lazy[v];
                lazy[2 * v + 1] %= MOD;
            }
 
            lazy[v] = 0;
        }
    }
 
    void range_update(int v, int tl, int tr, int l, int r, int value) {
        lazy_update(v, tl, tr);
 
        if (tr < l || tl > r) return;
 
        if (tl >= l and tr <= r) {
            lazy[v] += value;
            lazy[v] %= MOD;
            lazy_update(v, tl, tr);
        } else {
            int tm = tl + (tr - tl) / 2;
 
            range_update(2 * v, tl, tm, l, min(tm, r), value);
            range_update(2 * v + 1, tm + 1, tr, max(tm + 1, l), r, value);
            tree[v] = (tree[2 * v] + tree[2 * v + 1]) % MOD;
        }
    }
 
private:
    int n;
    vector<int> tree;
    vector<int> lazy;
};
 
class Solution {
public:
    int sumCounts(vector<int>& nums) {
        int N = nums.size();
        LazySegmentTree st(N);
        unordered_map<int, int> last_seen;
        long long curr = 0;
        long long res = 0;
 
        for (int i = 0; i < N; i++) {
            int x = nums[i];
            int last = (last_seen.find(x) != last_seen.end() ? last_seen[x]: -1) + 1;
            int queryLength = i - last + 1;
            curr = (curr + queryLength + 2 * st.query(1, 0, N - 1, last, i)) % MOD;
            st.range_update(1, 0, N - 1, last, i, 1);
            res = (res + curr) % MOD;
            last_seen[x] = i;
        }
 
        return (int) res;
    }
};