leetcode/0789_kth-largest-element-in-a-stream/python3/solution.py

27 lines
618 B
Python
Raw Permalink Normal View History

import heapq
class KthLargest:
def __init__(self, k: int, nums: List[int]):
self.k = k
self.heap = nums
heapq.heapify(self.heap)
# Keep removing elements until heap size is k
while len(self.heap) > k:
heapq.heappop(self.heap)
def add(self, val: int) -> int:
heapq.heappush(self.heap, val)
if len(self.heap) > self.k:
heapq.heappop(self.heap)
return self.heap[0]
# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)