aboutsummaryrefslogtreecommitdiff
path: root/subarr_distinct_elem_sum.py
blob: d66a2294b5486d45ec4f95a1891fdf826b5febe1 (plain)
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]))