Kth Largest Element in a Stream: Prioritizing Elements in a Real-time

Picture this: You’re listening to a live radio station that constantly plays songs. Every song has a popularity score. While listening, you want to know, at any moment, the 3rd most popular song played so far. This problem is not so different; it deals with finding the kth largest element from a stream of numbers.

Problem Statement

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. The class should have the following methods:

1. A constructor that takes an integer k and an integer array nums which will be the initial stream.
2. An add method that takes an integer and returns the element representing the kth largest element in the stream.

Brute Force Solution

A naive approach involves maintaining a list. Every time a new number is added, append it to the list, sort the list, and then get the kth largest element.

import java.util.*;

public class KthLargest {

    List<Integer> numsList;
    int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        this.numsList = new ArrayList<>();
        for (int num : nums) {
            this.numsList.add(num);
        }
    }

    public int add(int val) {
        numsList.add(val);
        Collections.sort(numsList);
        return numsList.get(numsList.size() - k);
    }
}

This approach is simple but inefficient. Sorting the list every time we add a number can become computationally heavy with long streams.

Optimal Solution

A more efficient way is to use a Priority Queue (or min-heap). Instead of sorting all numbers, we only track the top k numbers. When a new number is added, we check if it’s larger than the smallest of the top k numbers. If so, we remove the smallest and add the new number.

Imagine you have a transparent container that only shows the top k elements. Every time a new element is added, if it’s large enough, it pushes out the smallest of the top elements. If it’s smaller, it’s ignored. The challenge here is to efficiently manage this “transparent container.”

import java.util.*;

public class KthLargest {

    final PriorityQueue<Integer> q;
    final int k;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        q = new PriorityQueue<>(k);
        for (int n : nums) {
            add(n);
        }
    }

    public int add(int val) {
        if (q.size() < k) {
            q.offer(val);
        } else if (q.peek() < val) {
            q.poll();
            q.offer(val);
        }
        return q.peek();
    }
}

Explanation:

1. We maintain a priority queue (min-heap) with a maximum size of k. This heap only contains the top k largest elements seen so far.
2. When adding a new number:
   • If the heap size is less than k, we add the number.
   • If the heap size is equal to k, we compare the smallest number in the heap (q.peek()) with the new number. If the new number is larger, we remove the smallest number and add the new number.

This approach ensures that the heap always contains the top k numbers seen so far. The smallest number in the heap is always the kth largest overall.

Conclusion

The “Kth Largest Element in a Stream” problem is a classic example of problems where priority queues shine. It demonstrates the power of using the right data structure to simplify and optimize a solution. With the optimal approach, you’re not doing unnecessary work but only doing just enough to maintain the top k elements efficiently.