update readme

Author: hogan-tech

0315-count-of-smaller-numbers-after-self/0315-count-of-smaller-numbers-after-self.py

@@ -2,58 +2,66 @@
 # space complexity: O(n)
 from typing import List
 
-class TreeNode:
-    def __init__(self, start, end, val=0, left=None, right=None):
-        self.val = val
-        self.start = start
-        self.end = end
-        self.left = left
-        self.right = right
 
 class SegmentTree:
-    def __init__(self, n):
-        self.root = self.build(0, n - 1)
-
-    def build(self, l, r):
-        if l == r:
-            return TreeNode(l, r, 0)
-        leftTree = self.build(l, (l + r) // 2)
-        rightTree = self.build((l + r) // 2 + 1, r)
-        return TreeNode(l, r, 0, leftTree, rightTree)
-
-    def update(self, root, index, value):
-        if root.start == root.end == index:
-            root.val += value
-            return root.val
-        if root.start > index or root.end < index:
-            return root.val
-        root.val = self.update(root.left, index, value) + self.update(root.right, index, value)
-        return root.val
-
-    def query(self, root, l, r) -> int:
-        if root.start > r or root.end < l:
+    def __init__(self, size: int):
+        self.n = size
+        self.segTree = [0 for _ in range(4 * size)]
+
+    def update(self, idx: int, val: int, nodeIdx=0, start=0, end=None):
+        if end is None:
+            end = self.n - 1
+        if start == end:
+            self.segTree[nodeIdx] += val
+            return
+
+        mid = (start + end) // 2
+        leftIdx = 2 * nodeIdx + 1
+        rightIdx = 2 * nodeIdx + 2
+
+        if idx <= mid:
+            self.update(idx, val, leftIdx, start, mid)
+        else:
+            self.update(idx, val, rightIdx, mid + 1, end)
+
+        self.segTree[nodeIdx] = self.segTree[leftIdx] + self.segTree[rightIdx]
+
+    def query(self, left: int, right: int, nodeIdx=0, start=0, end=None) -> int:
+        if end is None:
+            end = self.n - 1
+        if right < start or left > end:
             return 0
-        if l <= root.start and root.end <= r:
-            return root.val
-        return self.query(root.left, l, r) + self.query(root.right, l, r)
+        if left <= start and end <= right:
+            return self.segTree[nodeIdx]
+
+        mid = (start + end) // 2
+        leftIdx = 2 * nodeIdx + 1
+        rightIdx = 2 * nodeIdx + 2
+        return self.query(left, right, leftIdx, start, mid) + self.query(left, right, rightIdx, mid + 1, end)
+
 
 class Solution:
     def countSmaller(self, nums: List[int]) -> List[int]:
         if not nums:
             return []
+
         sortedNums = sorted(set(nums))
         rankMap = {val: idx for idx, val in enumerate(sortedNums)}
+
         tree = SegmentTree(len(sortedNums))
         result = []
-        for n in reversed(nums):
-            idx = rankMap[n]
-            result.append(tree.query(tree.root, 0, idx - 1))
-            tree.update(tree.root, idx, 1)
+
+        for num in reversed(nums):
+            idx = rankMap[num]
+            result.append(tree.query(0, idx - 1))
+            tree.update(idx, 1)
+
         return result[::-1]
 
+
 nums = [5, 2, 6, 1]
-print(Solution().countSmaller(nums))
+print(Solution().countSmaller(nums))  
 nums = [-1]
-print(Solution().countSmaller(nums))
+print(Solution().countSmaller(nums))  
 nums = [-1, -1]
-print(Solution().countSmaller(nums))
+print(Solution().countSmaller(nums))