aboutsummaryrefslogtreecommitdiff
path: root/subarr_distinct_elem_sum.py
diff options
context:
space:
mode:
Diffstat (limited to 'subarr_distinct_elem_sum.py')
-rw-r--r--subarr_distinct_elem_sum.py68
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]))