Kth Largest Element in a Stream

Patrick Leaton
Published: March 16, 2021, 7:07 p.m.
Problem Description

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Implement KthLargest class:

  • KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of integers nums.
  • int add(int val) Returns the element representing the kth largest element in the stream.

 

Example 1:

Input
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
Output
[null, 4, 5, 5, 8, 8]

Explanation
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3);   // return 4
kthLargest.add(5);   // return 5
kthLargest.add(10);  // return 5
kthLargest.add(9);   // return 8
kthLargest.add(4);   // return 8

 

Constraints:

  • 1 <= k <= 10^4
  • 0 <= nums.length <= 10^4
  • -10^4 <= nums[i] <= 10^4
  • -10^4 <= val <= 10^4
  • At most 10^4 calls will be made to add.

 

The description was taken from https://leetcode.com/problems/kth-largest-element-in-a-stream/.

Problem Solution

#O(N(log K)) Time, O(K) Space
from heapq import heapify, heappush, heappop
class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.heap = nums
        heapify(self.heap)
 
    def add(self, val: int) -> int:
        heappush(self.heap, val)
        while len(self.heap) > self.k:
            heappop(self.heap)
        return self.heap[0]

Problem Explanation

If we're given a problem and are asked to return the kth largest or smallest element from a list, stream, string, etc., that is usually an indication that a heap may be used because at that point, the input will have to be sorted and if we are inserting values into a sorted list, it's better to utilize the O(Log(k)) insertion into a heap than into an array that would have to be sorted requiring O(N(Log(N))) time from each value having to be sorted again.

Python has a min-heap by default and no max-heap, so if we are trying to find the kth largest element from a min-heap there is a couple of tricks we could choose to do.

The first, and most common is to negate each value being pushed into the heap so that the largest value will become the smallest when negative.  We would then continuously pop an element and decrement k so that by the time k is zero, we would have the kth largest element. 

The second, and what we will choose to do here, is to use the length of the heap as the comparison value.  If we pop off each value until the length of the heap is k, then we would have popped off each value smaller than k and the kth largest element would be the smallest value in our heap which we would be able to access in constant time afterward.

Let's first initialize our k value and set the heap to our input list.

Then we will heapify the heap to sort it.

    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.heap = nums
        heapify(self.heap)

We’ll then make our add function that we will use to push values onto our heap.

    def add(self, val: int) -> int:
        heappush(self.heap, val)

 

After, we will use a loop that will run while the heap size is greater than our k value and it will pop the additional elements so we always have the kth largest element as the root once the loop finishes.

        while len(self.heap) > self.k:
            heappop(self.heap)

 

Since we have popped off k elements, the kth largest element will be the smallest element in our heap.

We can return this value from the top of the heap.

        return self.heap[0]