Kth Largest Element in a Stream

Problem Summary

• Given a list of nums and also a stream of incoming values. Define a class with __init__ and add() that returns the largest $k^{th}$ element

Approach Taken

• Use a min heap to always keep $k$ values and return $k^{th}$ largest element in $O(1)$
• This lets us achieve $O(logn)$ for our add()

Why This Works

• Heap has $O(logn)$ push or pop

Main Concepts Used

• Min Heap

Time & Space Complexity

• Time for init: $O(nlogn)$ → Reason: Initial heapify is $O(n)$, however reducing the size of the heap to $k$ is $logn$ and pops of $n-k$ times
• Time for add: $O(logn)$ → Reason: In the worst case, we will do both pop and push
• Space: $O(k)$ → Reason: We will only store $k$ items in our min heap

Code

class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.minHeap, self.k = nums, k
        heapq.heapify(self.minHeap)

        while len(self.minHeap) > k:
            heapq.heappop(self.minHeap)

    def add(self, val: int) -> int:
        if len(self.minHeap) < self.k:
            heapq.heappush(self.minHeap, val)
        elif val > self.minHeap[0]:
            heapq.heappush(self.minHeap, val)
            heapq.heappop(self.minHeap)
        return self.minHeap[0]

Common Mistakes / Things I Got Stuck On

• Not setting member variables correctly self.{variable}