diff options
Diffstat (limited to 'subarr_distinct_elem_sum.py')
-rw-r--r-- | subarr_distinct_elem_sum.py | 68 |
1 files changed, 68 insertions, 0 deletions
diff --git a/subarr_distinct_elem_sum.py b/subarr_distinct_elem_sum.py new file mode 100644 index 0000000..d66a229 --- /dev/null +++ b/subarr_distinct_elem_sum.py @@ -0,0 +1,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])) |