kth_lowest_with_updates.py

kth_lowest_with_updates.py

from typing import Generator
        
class TreeNode:
    def __init__(self, key, val):
        self.key = key
        self.val = val
        self.left = None
        self.right = None
        
    def __str__(self):
        return f'TreeNode (key: {self.key}, val: {self.val})\n'

class BST:
    def __init__(self):
        self.root = None

    def insert(self, node: TreeNode):
        def recur(root: TreeNode, insert: TreeNode) -> TreeNode:
            if root == None:
                return insert
            if root.val < insert.val:
                root.right = recur(root.right, insert)
            else:
                root.left = recur(root.left, insert)
                
            return root
        
        self.root = recur(self.root, node)
    
    def findGreaterThan(self, val: int) -> Generator[TreeNode, None, None]:
        def recur(root: TreeNode, val: int) -> Generator[TreeNode, None, None]:
            if root == None:
                return
            for node in recur(root.left, val):
                yield node
            if root.val > val:
                yield root
            for node in recur(root.right, val):
                yield node
                    
        for node in recur(self.root, val):
            yield node

    def find(self, val: int) -> TreeNode:
        def recur(root: TreeNode, val: int) -> TreeNode:
            if root == None:
                return
            if val < root.val:
                return recur(root.left, val)
            elif val == root.val:
                return root
            else:
                return recur(root.right, val)
            
        return recur(self.root, val)
    
    def delete(self, node: TreeNode) -> TreeNode:
        def recurS(root: TreeNode) -> TreeNode:
            if (root.left != None):
                return recurS(root.left)
            
            return root
        
        def recurD(root: TreeNode, node: TreeNode) -> TreeNode:
            if root == None:
                return
            if node.val < root.val:
                root.left = recurD(root.left, node)
            elif node.val > root.val:
                root.right = recurD(root.right, node)
            else:
                if root.left == None:
                    return root.right
                if root.right == None:
                    return root.left
                succ = recurS(root.right)
                root.key = succ.key
                root.val = succ.val
                root.right = recurD(root.right, succ)
            
            return root
        
        return recurD(self.root, node)
    
    def inorderTraversal(self):
        def recur(root: TreeNode):
            if root == None:
                return
            recur(root.left)
            print(root)
            recur(root.right)
        recur(self.root)

class Account:
    def __init__(self):
        self.bst = BST()
        self.insertions = 0
        self.lastInserted = None

    def insert(self, revenue: int):
        # update last inserted value in BST (re-insert it to not violate BST constraints)
        print(f'last_inserted is {self.lastInserted}')
        if self.lastInserted:
            self.bst.delete(self.lastInserted) # delete O(log n)
            self.bst.insert(TreeNode(self.lastInserted.key, self.lastInserted.val + revenue)) # update O(log n)
        self.lastInserted = TreeNode(self.insertions, revenue)
        self.bst.insert(self.lastInserted) # O(log n)
        self.insertions += 1

    def get_k_lowest(self, k: int, min_trevenue: int) -> Generator[int, None, None]:
        for n in self.bst.findGreaterThan(min_trevenue): # O(n)
            if (k == 0):
                return
            yield n.key
            k -= 1
            

if __name__ == '__main__':
    # bst = BST()
    # n1 = TreeNode(0, 5)
    # n2 = TreeNode(1, 2)
    # n3 = TreeNode(2, 4)
    # n4 = TreeNode(3, 7)
    # n5 = TreeNode(4, 9)
    # bst.insert(n1)
    # bst.insert(n2)
    # bst.insert(n3)
    # bst.delete(n4)
    # bst.insert(n4)
    # bst.insert(n5)
    # bst.inorderTraversal()
    # bst.delete(n2)
    # bst.inorderTraversal()
#     for n in bst.findGreaterThan(2):
#         print(n)
        
    ac = Account()
    ac.insert(8) # 0: 20
    ac.insert(12) # 1: 30
    ac.insert(18) # 2: 50
    ac.insert(32) # 3: 40
    ac.insert(8) # 4: 60
    ac.insert(52) # 5: 80
    ac.insert(28) # 6: 28
    
    for c in ac.get_k_lowest(4, 40):
        print(c)