leetcode

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

kth-largest-element-in-a-streamV5.py (2450B)

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