diff --git a/Count of Smaller Number before itself.py b/Count of Smaller Number before itself.py index b086b3b..cd44071 100644 --- a/Count of Smaller Number before itself.py +++ b/Count of Smaller Number before itself.py @@ -14,70 +14,92 @@ class Node(object): - def __init__(self, start, end): + def __init__(self, lo, hi): """Records the left subtree size""" - self.start = start - self.end = end + self.lo = lo + self.hi = hi self.left = None self.right = None self.cnt_this = 0 self.cnt_left = 0 def __repr__(self): - return repr("[%d, %d)" % (self.start, self.end)) + return repr("[%d, %d)" % (self.lo, self.hi)) class SegmentTree(object): def __init__(self): self.root = None - def build(self, root, start, end): - if start >= end: return - if not root: root = Node(start, end) + def build(self, root, lo, hi): + if lo >= hi: return + if not root: root = Node(lo, hi) - root.left = self.build(root.left, start, (end+start)/2) # not (end-start)/2 - if root.left: root.right = self.build(root.right, (end+start)/2, end) + root.left = self.build(root.left, lo, (lo+hi)/2) + if root.left: root.right = self.build(root.right, (lo+hi)/2, hi) return root def set(self, root, i, val): - if root.start == i and root.end-1 == root.start: + if root.lo == i and root.hi-1 == root.lo: root.cnt_this += val - elif i < (root.end+root.start)/2: + elif i < (root.lo+root.hi)/2: root.cnt_left += val self.set(root.left, i, val) else: self.set(root.right, i, val) def get(self, root, i): - if root.start == i and root.end-1 == root.start: + if root.lo == i and root.hi-1 == root.lo: return root.cnt_left - elif i < (root.end+root.start)/2: + elif i < (root.lo+root.hi)/2: return self.get(root.left, i) else: return root.cnt_left+root.cnt_this+self.get(root.right, i) class Solution(object): + def _build_tree(self, A): + st = SegmentTree() + mini, maxa = min(A), max(A) + st.root = st.build(st.root, mini, maxa+2) # maxa+1 is the end dummy + return st + def countOfSmallerNumberII(self, A): """ + count of smaller number before it. Segment Tree :param A: A list of integer :return: Count the number of element before this element 'ai' is smaller than it and return count number list """ if not A: return [] - st = SegmentTree() - st.root = st.build(st.root, min(A), max(A)+1) + st = self._build_tree(A) + ret = [] + for a in A: + st.set(st.root, a, 1) + ret.append( + st.get(st.root, a) + ) + + return ret + + def countOfLargerElementsBeforeElement(self, A): + if not A: return [] + st = self._build_tree(A) ret = [] + end = max(A)+1 for a in A: - ret.append(st.get(st.root, a)) + ret.append( + st.get(st.root, end) - st.get(st.root, a) + ) st.set(st.root, a, 1) return ret if __name__ == "__main__": - print Solution().countOfSmallerNumberII([1, 2, 7, 8, 5]) + assert Solution().countOfSmallerNumberII([1, 2, 7, 8, 5]) == [0, 1, 2, 3, 2] + assert Solution().countOfLargerElementsBeforeElement([1, 9, 7, 8, 5]) == [0, 0, 1, 1, 3] assert Solution().countOfSmallerNumberII( [26, 78, 27, 100, 33, 67, 90, 23, 66, 5, 38, 7, 35, 23, 52, 22, 83, 51, 98, 69, 81, 32, 78, 28, 94, 13, 2, 97, 3, 76, 99, 51, 9, 21, 84, 66, 65, 36, 100, 41]) == [0, 1, 1, 3, 2, 3, 5, 0, 4, 0, 5, 1, 6, 2, 9, 2, 14, 10, 17,