leetcode

Leetcode submissions
git clone git://git.laack.co/leetcode.git
Log | Files | Refs | README

kth-largest-element-in-a-streamV3.py (2093B)

      1 # TODO:
      2 # I still need to implement heappush and heappop.
      3 
      4 from heapq import heappush, heappop
      5 
      6 
      7 class Heap:
      8     def __init__(self, nums):
      9         if nums is None:
     10             self.nums = []
     11             return
     12 
     13         if len(nums) <= 1:
     14             self.nums = nums
     15             return
     16 
     17         self.nums = self._make_heap(nums, 0)
     18 
     19     def _make_heap(self, nums, index):
     20         if not self._is_leaf(nums, index):
     21             self._make_heap(nums, self._left_child(index))
     22             self._make_heap(nums, self._right_child(index))
     23         else:
     24             return
     25 
     26         self._compare_three(nums, index)
     27         return nums
     28 
     29     def _left_child(self, index):
     30         return (index * 2) + 1
     31 
     32     def _right_child(self, index):
     33         return (index * 2) + 2
     34 
     35     def pop(self):
     36         return heappop(self.nums)
     37 
     38     def push(self, val):
     39         heappush(self.nums, val)
     40 
     41     def _compare_three(self, nums, index):
     42         """Compare and move pair of three to be arranged correctly. Assume index is root."""
     43 
     44         root = index
     45         left = self._left_child(index)
     46         right = self._right_child(index)
     47 
     48         min_index = root
     49 
     50         if nums[left] < nums[min_index]:
     51             min_index = left
     52         if len(nums) > right and nums[right] < nums[min_index]:
     53             min_index = right
     54 
     55         if min_index != index:
     56             temp = nums[root]
     57             nums[root] = nums[min_index]
     58             nums[min_index] = temp
     59             if not self._is_leaf(nums, min_index):
     60                 self._compare_three(nums, min_index)
     61 
     62         return
     63 
     64     def _is_leaf(self, nums, index):
     65 
     66         if len(nums) <= (index * 2) + 1:
     67             return True
     68 
     69         return False
     70 
     71 
     72 class KthLargest:
     73 
     74     def __init__(self, k: int, nums: List[int]):
     75 
     76         heap = Heap(nums)
     77         self.heap = heap
     78         self.k = k
     79         self.prune()
     80 
     81     def prune(self):
     82         k = self.k
     83         heap = self.heap
     84         while len(heap.nums) > k:
     85             heap.pop()
     86 
     87     def add(self, val: int) -> int:
     88         self.heap.push(val)
     89         self.prune()
     90         return self.heap.nums[0]