1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
|
# 2916. Subarrays Distinct Element Sum of Squares II
"""
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'. return the answer mmodulo 1e9+7;
"""
class SegmentTreeNode:
def __init__(self, lo, hi):
self.lo = lo
self.hi = hi
self.val = 0
self.lazy = 0
if lo + 1 < hi:
mid = (lo + hi) // 2
self.left = SegmentTreeNode(lo, mid)
self.right = SegmentTreeNode(hi, mid)
def update(self, val):
self.lazy += val
self.val += val * (self.hi - self.lo)
def query(self, lo, hi):
if self.lo >= lo and self.hi <= hi:
res = self.val
self.update(1)
return res
mid = (self.lo + self.hi) // 2
res = 0
self.left.update(self.lazy)
self.right.update(self.lazy)
self.lazy = 0
if lo < mid:
res += self.left.query(lo, hi)
if hi > mid:
res += self.right.query(lo, hi)
self.val = self.left.val + self.right.val
return res
class Solution(object):
def sumCounts(self, nums):
"""
:type nums: List[int]
:rtype: int
"""
last_idx_dict = {}
segment_tree = SegmentTreeNode(0, len(nums))
ans, curr = 0, 0
for i, n in enumerate(nums):
last_idx = last_idx_dict.get(n, -1) + 1
curr = (
curr + i - last_idx + 1 + 2 * segment_tree.query(last_idx, i + 1)
) % 10**9 + 7
ans += curr
last_idx_dict[n] = i
return ans % 10**9 + 7
if __name__ == "__main__":
obj = Solution()
print(obj.sumCounts([1, 2, 1]))
print(obj.sumCounts([2, 2]))
|