Kth Largest Element in a Stream — Min-Heap Design [Amazon Easy]

Sanjeev SharmaSanjeev Sharma
2 min read

Advertisement

Problem Statement

Design a class that finds the kth largest element in a stream. Implement KthLargest(k, nums) and int add(val).

Example: k=3, nums=[4,5,8,2], add(3)→4, add(5)→5, add(10)→8

Intuition

Keep a min-heap of exactly k elements. The heap top (minimum of the top-k) is the kth largest. When adding: push, then pop if size > k.

C++ Solution

#include <queue>
#include <vector>
using namespace std;
class KthLargest {
    priority_queue<int, vector<int>, greater<int>> pq;
    int k;
public:
    KthLargest(int k, vector<int>& nums) : k(k) {
        for (int n : nums) add(n);
    }
    int add(int val) {
        pq.push(val);
        if ((int)pq.size() > k) pq.pop();
        return pq.top();
    }
};

Java Solution

import java.util.PriorityQueue;
class KthLargest {
    PriorityQueue<Integer> pq = new PriorityQueue<>();
    int k;
    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int n : nums) add(n);
    }
    public int add(int val) {
        pq.offer(val);
        if (pq.size() > k) pq.poll();
        return pq.peek();
    }
}

Python Solution

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]

JavaScript Solution

// Use a sorted array (or MinHeap library) — simplified approach
class KthLargest {
    constructor(k, nums) {
        this.k = k;
        this.data = nums.sort((a,b)=>a-b).slice(-k);
    }
    add(val) {
        // Binary insert then trim
        let lo=0, hi=this.data.length;
        while(lo<hi){const m=(lo+hi)>>1; this.data[m]<val?lo=m+1:hi=m;}
        this.data.splice(lo,0,val);
        if(this.data.length>this.k) this.data.shift();
        return this.data[0];
    }
}

Complexity: add → O(log k) time, O(k) space

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro