Kth Largest Element in a Stream

Sanjeev SharmaSanjeev Sharma
2 min read

Advertisement

Problem

Design a class that adds numbers from a stream and returns the Kth largest element after each addition.

Key insight: Maintain a min-heap of size K. The minimum of the heap is the Kth largest.

Approach

  • Init: heapify first K elements, push remaining one-by-one evicting if > K
  • Add: push to heap, pop if size > K, return heap[0]

Solutions

// C — conceptual (use array-based min-heap)
// See C++ for clean implementation
// C++
class KthLargest {
    priority_queue<int, vector<int>, greater<int>> minH;
    int k;
public:
    KthLargest(int k, vector<int>& nums) : k(k) {
        for (int n : nums) add(n);
    }
    int add(int val) {
        minH.push(val);
        if ((int)minH.size() > k) minH.pop();
        return minH.top();
    }
};
// Java
class KthLargest {
    PriorityQueue<Integer> minH = new PriorityQueue<>();
    int k;
    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int n : nums) add(n);
    }
    public int add(int val) {
        minH.offer(val);
        if (minH.size() > k) minH.poll();
        return minH.peek();
    }
}
// JavaScript (using sorted array for simplicity)
class KthLargest {
    constructor(k, nums) {
        this.k = k;
        this.heap = [];
        for (const n of nums) this.add(n);
    }
    add(val) {
        this.heap.push(val);
        this.heap.sort((a, b) => a - b);
        if (this.heap.length > this.k) this.heap.shift();
        return this.heap[0];
    }
}
# Python
import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.heap = []
        for n in nums:
            self.add(n)

    def add(self, val):
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

Complexity

  • add(): O(log k)
  • Space: O(k)

Key Insight

A min-heap of size K always has the Kth largest at the top (the minimum). Elements smaller than the Kth largest are evicted.

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro