Solution: Kth Largest Element in a Stream

Let's solve the Kth Largest Element in a Stream problem using the Top K Elements pattern.

Statement

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

Note: It is the kthk^{th} largest element in the sorted order, not the kthk^{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 kthk^{th} largest element in the stream.

Constraints:

  • 1k1031 \leq k \leq 10^3
  • 00 \leq nums.length 103\leq 10^3
  • 103-10^3 \leq nums[i] 103\leq 10^3
  • 103-10^3 \leq value 103\leq 10^3
  • At most, 10310^3 calls will be made to add.
  • It is guaranteed that there will be at least kk elements in the array when you search for the kthk^{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 kthk^{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(n2)O(n^2), where nn is the number of elements in the data stream. The time complexity of each insertion is O(n)O(n) and finding the kthk^{th} largest element would take O(1)O(1) time, assuming we are storing the data in an array. The space complexity is O(1)O(1).

Optimized approach using Top K Elements

As new elements are added to the number stream, the kthk^{th} largest element keeps changing. We need to implement a class that caters to the dynamically changing numbers. The most efficient data structure for repeatedly finding the kthk^{th} largest number in a changing list is a heap.

We’ll implement a min-heap of size kk. In a min-heap, the smallest number is always at the top. We’ll use this property to design a solution that ensures that in a min-heap with kk elements, the kthk^{th} largest element is always at the top of the heap.

The slides below illustrate the core ideas of our algorithm:

Level up your interview prep. Join Educative to access 70+ hands-on prep courses.