python | 100% faster | bst solution explained
So the challenge here is how can we quickly sum the k-largest and k-smallest elements, and I want to talk about how I did that with a binary search tree. The ke
So the challenge here is how can we quickly sum the k-largest and k-smallest elements, and I want to talk about how I did that with a binary search tree. The key idea was to store the number of items in each sub-tree, as well as the sum of all the items in each sub tree. By storing the size of each node it becomes possible to determine if the ith smallest node is to the left of the current node, at the current node, or to the right of the current node. This means the ith smallest node can be found in a single trip down the tree. Because the sums are stored, we can also choose to sum elements which are smaller or larger than the target as we go (without having to actually traverse the tree).\n\nThe other aspects of my algorithm are fairly straightforward. I used a deque to store the last m elements in the stream, and an integer to store the sliding window sum for all of the last m elements. This means that to compute the average I have to subtract the sum of the k largest/smallest elements from the sliding window sum, and then divide by the number of elements (which is pre-computed since it\'s a constant). Whenever I add/remove an element to/from the queue I also/remove add it to/from the tree and update the sum, with different cases to handle when the tree is null, when the queue isn\'t full.\n\nBST Height (n is the number of nodes in the tree): \nBest : O(logn)\nAverage : O(logn)\nWorst : O(n)\n\nBST Runtimes (h is the height of the tree):\n\\_\\_init\\_\\_ : O(1)\ninsert/remove : O(h)\nklargets/ksmallest : O(h)\n\nMKAverage Runtimes:\n\\_\\_init\\_\\_ : O(1)\naddElement : O(h)\ncalculateMKAverage : O(h)\n\nBest Performance: 864 ms/49.3 MB : 100%/98%\n\nA more complete version of the binary search tree is below the code for this puzzle.\n\'\'\'\n\n\t#Binary Search Tree\n\tclass Node:\n\t\t#Create a Node\n\t\tdef __init__(self, val:int) -> None:\n\t\t\t"""Create a New Binary Search Tree Node"""\n\t\t\tself.val = val\n\t\t\tself.sum = val\n\t\t\tself.size = 1\n\t\t\tself.count = 1\n\t\t\tself.left = None\n\t\t\tself.right = None\n\t\t\n\t\t#Add a Node to the Tree\n\t\tdef insert(self, val:int) -> None:\n\t\t\t"""Add a Value to the Tree"""\n\t\t\t#Update the Size/Sum\n\t\t\tself.size += 1\n\t\t\tself.sum += val\n\t\t\t\n\t\t\t#Check the Case\n\t\t\tif val < self.val:\n\t\t\t\t#Check the Left Node\n\t\t\t\tif self.left:\n\t\t\t\t\t#Recurse\n\t\t\t\t\tself.left.insert(val)\n\t\t\t\telse:\n\t\t\t\t\t#Make a New Node\n\t\t\t\t\tself.left = Node(val)\n\t\t\telif val > self.val:\n\t\t\t\t#Check the Right Node\n\t\t\t\tif self.right:\n\t\t\t\t\t#Recurse\n\t\t\t\t\tself.right.insert(val)\n\t\t\t\telse:\n\t\t\t\t\t#Make a New Node\n\t\t\t\t\tself.right = Node(val)\n\t\t\telse:\n\t\t\t\t#Increment the Count\n\t\t\t\tself.count += 1\n\n\t\t#Remove a Node from the Tree (Safetys Removed)\n\t\tdef remove(self, val:int) -> None:\n\t\t\t"""Remove a Value From the Tree"""\n\t\t\t#Update the Size/Sum\n\t\t\tself.size -= 1\n\t\t\tself.sum -= val\n\t\t\t\n\t\t\t#Check the Case\n\t\t\tif val < self.val:\n\t\t\t\t#Recurse Left\n\t\t\t\treturn self.left.remove(val)\n\t\t\telif val > self.val:\n\t\t\t\t#Recurse Right\n\t\t\t\treturn self.right.remove(val)\n\t\t\telse:\n\t\t\t\t#Decrement the Count\n\t\t\t\tself.count -= 1\n\n\t\t#Sum the k Smallest Numbers\n\t\tdef ksmallest(self, k:int) -> int:\n\t\t\t"""Efficiently Sum the k Smallest Values in the Tree"""\n\t\t\t#Check the Left Branch\n\t\t\ts = 0\n\t\t\tif self.left:\n\t\t\t\t#Check the Case\n\t\t\t\tif k <= self.left.size:\n\t\t\t\t\t#Recurse Left\n\t\t\t\t\treturn self.left.ksmallest(k)\n\t\t\t\telse:\n\t\t\t\t\t#Update k/s\n\t\t\t\t\tk -= self.left.size\n\t\t\t\t\ts = self.left.sum\n\n\t\t\t#Check the Current Value\n\t\t\tif k <= self.count:\n\t\t\t\treturn s + k*self.val\n\t\t\telif self.right:\n\t\t\t\t#Recurse Right\n\t\t\t\treturn s + self.count*self.val + self.right.ksmallest(k - self.count)\n\t\t\telse:\n\t\t\t\t#Return Search Failure\n\t\t\t\treturn None\n\n\tclass MKAverage:\n\t\tdef __init__(self, m: int, k: int):\n\t\t\tself.k = k\n\t\t\tself.l = m - k\n\t\t\tself.m = m\n\t\t\tself.n = self.m - 2*self.k\n\t\t\tself.window = deque()\n\t\t\tself.tree = None\n\n\t\tdef addElement(self, num: int) -> None:\n\t\t\t#Check the Case\n\t\t\tif len(self.window) == self.m:\n\t\t\t\t#Remove the Previous Number from the Window/Tree\n\t\t\t\tself.tree.remove(self.window.popleft())\n\n\t\t\t\t#Add the New Number to the Window/Tree\n\t\t\t\tself.window.append(num)\n\t\t\t\tself.tree.insert(num)\n\t\t\telif self.window:\n\t\t\t\t#Add the New Number to the Window/Tree\n\t\t\t\tself.window.append(num)\n\t\t\t\tself.tree.insert(num)\n\t\t\telse:\n\t\t\t\t#Make a New Tree\n\t\t\t\tself.window.append(num)\n\t\t\t\tself.tree = Node(num)\n\n\t\tdef calculateMKAverage(self) -> int:\n\t\t\t#Check the Window\n\t\t\tif len(self.window) < self.m:\n\t\t\t\treturn -1\n\t\t\telse:\n\t\t\t\t#Compute Return the Average\n\t\t\t\treturn (self.tree.ksmallest(self.l) - self.tree.ksmallest(self.k))//self.n\n\'\'\'\n\nHere is my binary search tree class as it exists on my computer!\n\n\'\'\'\n\n    #Binary Search Tree\n    class Node:\n        #Create a Node\n        def __init__(self, val:int) -> None:\n            """Create a New Binary Search Tree Node"""\n            self.val = val\n            self.sum = val\n            self.size = 1\n            self.count = 1\n            self.left = None\n            self.right = None\n        \n        #Get the Size of the Tree\n        def __len__(self) -> int:\n            """Return the Number of Items in the Tree"""\n            return self.size\n        \n        #Search for an Item (Index Validity Guaranteed)\n        def find(self, i:int):\n            """Find and Return the Node at Index i\n            \n            It also returns an index describing which version of the \n            current node was returned.\n            \n            If i is not a valid index then it returns None (but still \n            does a non-trivial amount of work to figure that out, so \n            be cautious when using this)\n            """\n            #Check the Left Branch\n            if self.left:\n                #Check the Case\n                if i < self.left.size:\n                    #Recurse Left\n                    return self.left.find(i)\n                else:\n                    #Update i\n                    i -= self.left.size\n            \n            #Check the Current Value\n            if i < self.count:\n                return self, i\n            elif self.right:\n                #Recurse Right\n                return self.right.find(i - self.count)\n            else:\n                #Return Search Failure\n                return None\n        \n        #Find the ith Smallest Element\n        def __getitem__(self, i:int) -> int:\n            """Get and Return the Value of the Node at Index i"""\n            #Check i\n            if i < -self.size or i >= self.size:\n                #Index Out of Range\n                return None\n            elif i >= 0:\n                #Find the Index Recursively\n                return self.find(i)[0].val\n            else:\n                #Find the Index Recursively\n                return self.find(self.size + i)[0].val\n        \n        #Count the Number of Occurences of an Item in the Tree\n        def Count(self, val:int) -> int:\n            #Check the Case\n            if val == self.val:\n                #Return the Count\n                return self.count\n            elif val < self.val:\n                #Recurse Left\n                return self.left.Count(val) if self.left else 0\n            else:\n                #Recurse Right\n                return self.right.Count(val) if self.right else 0\n        \n        #Check if an Item is in the Tree\n        def __contains__(self, val:int) -> bool:\n            """Check Whether a Value is in the Tree"""\n            return self.Count(val) > 0\n        \n        #Find the Minimum Item\n        def getMin(self) -> int:\n            """Returns the Minimum Value in the Tree"""\n            return self[0]\n        \n        #Find the Maximum Item\n        def getMax(self) -> int:\n            """Returns the Maximum Valud in the Tree"""\n            return self[-1]\n        \n        #Find the Median\n        def median(self) -> float:\n            """Returns the Median Value in the Tree"""\n            #Check the Case\n            if self.size%2:\n                #The Median is a Single Item\n                return self[self.size//2]\n            else:\n                #Find the Larger Item\n                node, i = self.find(self.size//2)\n                \n                #Check the Case\n                if i > 0:\n                    #The Median is Already Found\n                    return node.val\n                elif node.left and node.left.size:\n                    #Get the Largest Value from the Left\n                    return (node.val + node.left[-1])/2\n                else:\n                    #Do a Separate Search\n                    return (node.val + self[self.size//2 - 1])/2\n        \n        #Find the Average\n        def mean(self) -> float:\n            """Returns the Average Value in the Tree"""\n            return self.sum/self.size\n        \n        #Add a Node to the Tree\n        def insert(self, val:int) -> None:\n            """Add a Value to the Tree"""\n            #Check the Case\n            self.size += 1\n            self.sum += val\n            if val < self.val:\n                #Check the Left Node\n                if self.left:\n                    #Recurse\n                    self.left.insert(val)\n                else:\n                    #Make a New Node\n                    self.left = Node(val)\n            elif val > self.val:\n                #Check the Right Node\n                if self.right:\n                    #Recurse\n                    self.right.insert(val)\n                else:\n                    #Make a New Node\n                    self.right = Node(val)\n            else:\n                #Increment the Count\n                self.count += 1\n        \n        #Remove a Node from the Tree\n        def remove(self, val:int) -> bool:\n            """Remove a Value From the Tree"""\n            #Check the Case\n            if val < self.val:\n                #Check the Left Node\n                if self.left and self.left.remove(val):\n                    #Decrement the Sum/Size\n                    self.sum -= val\n                    self.size -= 1\n                    \n                    #Return Success\n                    return True\n                else:\n                    #The Node Wasn\'t Found\n                    return False\n            elif val > self.val:\n                #Check the Right Node\n                if self.right and self.right.remove(val):\n                    #Decrement the Sum/Size\n                    self.sum -= val\n                    self.size -= 1\n                    \n                    #Return Success\n                    return True\n                else:\n                    #The Node Wasn\'t Found\n                    return False\n            else:\n                #Check the Count\n                if self.count:\n                    #Decrement the Count, Sum, and Size\n                    self.count -= 1\n                    self.sum -= val\n                    self.size -= 1\n                    \n                    #Return Success\n                    return True\n                else:\n                    #Return Failure\n                    return False\n        \n        #Pop an Item at a Specified Index\n        def pop(self, i:int) -> int:\n            """Remove and Return the Item at Index i"""\n            #Check i\n            if i < 0 or i >= self.size:\n                #Return Search Failure\n                return None\n            \n            #Update the Size\n            self.size -= 1\n            \n            #Check the Left Branch\n            if self.left:\n                #Check the Case\n                if i < self.left.size:\n                    #Recurse Left\n                    val = self.left.pop(i)\n                    \n                    #Update the Sum\n                    self.sum -= val\n                    \n                    #Return the Value\n                    return val\n                else:\n                    #Update i\n                    i -= self.left.size\n            \n            #Check the Current Value\n            if i < self.count:\n                #Update the Count/Sum\n                self.count -= 1\n                self.sum -= self.val\n                \n                #Return the Value\n                return self.val\n            else:\n                #Recurse Right\n                val = self.right.pop(i - self.count)\n                \n                #Update the Sum\n                self.sum -= val\n                \n                #Return the Value\n                return val\n        \n        #Yield Nodes in Increasing Order\n        def increasing(self, unique:bool = False) -> int:\n            """Yield the Values of the Tree in Increasing Order"""\n            #Check the Left Node\n            if self.left:\n                #Recurse\n                yield from self.left.increasing(unique)\n            \n            #Yield the Current Node\n            if unique:\n                yield self.val, self.count\n            else:\n                yield from (self.val for i in range(self.count))\n            \n            #Check the Right Node\n            if self.right:\n                #Recurse\n                yield from self.right.increasing(unique)\n        \n        #Make the Tree Iterable\n        def __iter__(self):\n            return self.increasing()\n        \n        #Yield Nodes in Decreasing Order\n        def decreasing(self, unique:bool = False) -> int:\n            """Yield the Values in the Tree in Decreasing Order"""\n            #Check the Right Node\n            if self.right:\n                #Recurse\n                yield from self.right.decreasing(unique)\n            \n            #Yield the Current Node\n            if unique:\n                yield self.val, self.count\n            else:\n                yield from (self.val for i in range(self.count))\n            \n            #Check the Left Node\n            if self.left:\n                #Recurse\n                yield from self.left.decreasing(unique)\n        \n        #Sum the k Largest Numbers\n        def klargest(self, k:int) -> int:\n            """Efficiently Sum the k Largest Values in the Tree"""\n            #Check the Right Branch\n            s = 0\n            if self.right:\n                #Check the Case\n                if k <= self.right.size:\n                    #Recurse Left\n                    return self.right.klargest(k)\n                else:\n                    #Update k/s\n                    k -= self.right.size\n                    s = self.right.sum\n            \n            #Check the Current Value\n            if k <= self.count:\n                return s + k*self.val\n            elif self.left:\n                #Recurse Left\n                return s + self.count*self.val + self.left.klargest(k - self.count)\n            else:\n                #Return Search Failure\n                return None\n        \n        #Sum the k Smallest Numbers\n        def ksmallest(self, k:int) -> int:\n            """Efficiently Sum the k Smallest Values in the Tree"""\n            #Check the Left Branch\n            s = 0\n            if self.left:\n                #Check the Case\n                if k <= self.left.size:\n                    #Recurse Left\n                    return self.left.ksmallest(k)\n                else:\n                    #Update k/s\n                    k -= self.left.size\n                    s = self.left.sum\n            \n            #Check the Current Value\n            if k <= self.count:\n                return s + k*self.val\n            elif self.right:\n                #Recurse Right\n                return s + self.count*self.val + self.right.ksmallest(k - self.count)\n            else:\n                #Return Search Failure\n                return None\n\'\'\'
['Binary Search Tree', 'Python']