<p>You are given a <strong>0-indexed </strong>integer array <code>nums</code>.</p>

<p>The <strong>distinct count</strong> of a subarray of <code>nums</code> is defined as:</p>

<ul>
	<li>Let <code>nums[i..j]</code> be a subarray of <code>nums</code> consisting of all the indices from <code>i</code> to <code>j</code> such that <code>0 &lt;= i &lt;= j &lt; nums.length</code>. Then the number of distinct values in <code>nums[i..j]</code> is called the distinct count of <code>nums[i..j]</code>.</li>
</ul>

<p>Return <em>the sum of the <strong>squares</strong> of <strong>distinct counts</strong> of all subarrays of </em><code>nums</code>.</p>

<p>Since the answer may be very large, return it <strong>modulo</strong> <code>10<sup>9</sup> + 7</code>.</p>

<p>A subarray is a contiguous <strong>non-empty</strong> sequence of elements within an array.</p>

<p>&nbsp;</p>
<p><strong>Example 1:</strong></p>

<pre>
<strong>Input:</strong> nums = [1,2,1]
<strong>Output:</strong> 15
<strong>Explanation:</strong> 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 1<sup>2</sup> + 1<sup>2</sup> + 1<sup>2</sup> + 2<sup>2</sup> + 2<sup>2</sup> + 2<sup>2</sup> = 15.
</pre>

<p><strong>Example 2:</strong></p>

<pre>
<strong>Input:</strong> nums = [2,2]
<strong>Output:</strong> 3
<strong>Explanation:</strong> 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 1<sup>2</sup> + 1<sup>2</sup> + 1<sup>2</sup> = 3.</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= nums.length &lt;= 10<sup>5</sup></code></li>
	<li><code>1 &lt;= nums[i] &lt;= 10<sup>5</sup></code></li>
</ul>

from collections import defaultdict
class Solution:
    def sumCounts(self, nums: List[int]) -> int:
        n = len(nums)
        tree = [0]*(4*n)
        lazy = [0]*(4*n)
        def build(tree, arr, v, lo, hi):
            if lo == hi:
                tree[v] = arr[lo]
                return
            m = (lo + hi)//2
            build(tree, arr, 2*v + 1, lo, m)
            build(tree, arr, 2*v + 2, m + 1, hi)
            tree[v] = tree[2*v + 1] + tree[2*v + 2]

        def update(tree, v, lo, hi, i, j, val):
            if lazy[v]:
                tree[v] += (hi - lo + 1) * lazy[v]
                if lo < hi:
                    lazy[2*v + 1] += lazy[v]
                    lazy[2*v + 2] += lazy[v]
                lazy[v] = 0
            if lo > j or hi < i: 
                return 
            if i <= lo and hi <= j:
                tree[v] += (hi - lo + 1)*val
                if lo < hi:
                    lazy[2*v + 1] += val 
                    lazy[2*v + 2] += val
                return
            m = (lo + hi)//2
            update(tree, 2*v + 1, lo, m, i, j, val) 
            update(tree, 2*v + 2, m+1, hi, i, j, val)
            tree[v] = tree[2*v + 1] + tree[2*v + 2]

        def query(tree, v, lo, hi, i, j):
            if j < lo or hi < i:
                return 0
            if lazy[v]:
                tree[v] += (hi - lo + 1)*lazy[v]
                if lo < hi:
                    lazy[2*v + 1] += lazy[v]
                    lazy[2*v + 2] += lazy[v]
                lazy[v] = 0
            if i <= lo and hi <= j:
                return tree[v]
            m = (lo + hi)//2
            if i > m:
                return query(tree, 2*v + 2, m + 1, hi, i, j)
            elif j <= m:
                return query(tree, 2*v + 1, lo, m, i, j)
            return query(tree, 2*v + 1, lo, m, i, m) + query(tree, 2*v + 2, m+1, hi, m+1, j)

        squares = 1
        update(tree, 0, 0, n-1, n-1, n-1, 1)
        index = defaultdict(lambda:n)
        index[nums[-1]] = n - 1
        res = 1
        for i in range(n-2,-1,-1):
            j = index[nums[i]]
            squares += (j-i) + 2*query(tree,0,0,n-1,i+1,j-1)
            res += squares
            update(tree, 0, 0, n-1, i, j-1, 1)
            index[nums[i]] = i
        return res%(10**9 + 7)
6384