Solution: Kth Largest Element in a Stream

Statement

Given an infinite stream of integers (sorted or unsorted), nums, design a class to find the $k^{th}$ largest element in a stream.

Note: It is the $k^{th}$ largest element in the sorted order, not the $k^{th}$ distinct element.

The class should have the following functions, inputs, and return values:

• Init(nums, k): It takes an array of integers nums and an integer k and initializes the class object.

• Add(value): It takes one integer value, appends it to the stream, and returns the element representing the $k^{th}$ largest element in the stream.

Constraints:

• $1 \leq k \leq 10^3$
• $0 \leq$ nums.length $\leq 10^3$
• $-10^3 \leq$ nums[i] $\leq 10^3$
• $-10^3 \leq$ value $\leq 10^3$
• At most, $10^3$ calls will be made to add.
• It is guaranteed that there will be at least $k$ elements in the array when you search for the $k^{th}$ element.

Solution

So far, you have probably brainstormed some approaches and have an idea of how to solve this problem. Let’s explore some of these approaches and figure out which one to follow based on considerations such as time complexity and any implementation constraints.

Naive approach

The naive solution is to first sort the data and then find the $k^{th}$ largest element. Insertion sort is an algorithm that can be used to sort the data as it appears. However, it also requires shifting the elements, greater than the inserted number, one place forward.

The overall time complexity of the algorithm becomes $O(n^2)$, where $n$ is the number of elements in the data stream. The time complexity of each insertion is $O(n)$ and finding the $k^{th}$ largest element would take $O(1)$ time, assuming we are storing the data in an array. The space complexity is $O(1)$.

Optimized approach using Top K Elements

To efficiently find the $k^{th}$ largest element in a stream of numbers, we use a min-heap that holds the top $k$ largest elements. This way, we don’t have to sort the entire list each time a new number is added. The $k^{th}$ largest element will change as new members come in, so we need a class to handle these dynamic updates.

With its ability to hold k elements, the min-heap ensures that the $k^{th}$ largest number is always at the root. We do this by adding new elements to the heap and removing the smallest one if the heap grows beyond k elements. This approach allows us quick access to the $k^{th}$ largest number, making the min-heap the most efficient tool for the job.

The slides below illustrate the core ideas of our algorithm: