leetcode

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

kth-largest-element-in-a-streamV4.py (2242B)

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