Amazon — Kth Largest Element in a Stream (Min-Heap)

Problem (Amazon Stream Classic)

Design a class that finds the kth largest element in a stream:

KthLargest(k, nums) — initialize with k and initial array
add(val) — add value and return kth largest

Key Insight — Min-Heap of Size K

Maintain a min-heap of exactly k elements. The minimum (top) is always the kth largest.

When adding: push value, if size > k, pop the min. Top is answer.

Solutions

Python

import heapq

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

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

JavaScript

class KthLargest {
    constructor(k, nums) { this.k=k; this.h=[]; for(const n of nums) this.add(n); }
    add(val) {
        this._push(val);
        if(this.h.length>this.k) this._pop();
        return this.h[0];
    }
    _push(v){ this.h.push(v); let i=this.h.length-1; while(i>0){const p=(i-1)>>1;if(this.h[p]>this.h[i]){[this.h[p],this.h[i]]=[this.h[i],this.h[p]];i=p;}else break;}}
    _pop(){ [this.h[0],this.h[this.h.length-1]]=[this.h[this.h.length-1],this.h[0]]; this.h.pop(); let i=0; while(true){let l=2*i+1,r=2*i+2,m=i;if(l<this.h.length&&this.h[l]<this.h[m])m=l;if(r<this.h.length&&this.h[r]<this.h[m])m=r;if(m===i)break;[this.h[i],this.h[m]]=[this.h[m],this.h[i]];i=m;}}
}

Java

import java.util.*;
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); while(pq.size()>k) pq.poll(); return pq.peek(); }
}

C++

#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); while((int)pq.size()>k) pq.pop(); return pq.top(); }
};

C

int heap[100001]; int sz=0, K;
void push(int v){ heap[sz++]=v; int i=sz-1; while(i>0){int p=(i-1)/2;if(heap[p]>heap[i]){int t=heap[p];heap[p]=heap[i];heap[i]=t;i=p;}else break;} }
void pop(){ heap[0]=heap[--sz]; int i=0; while(1){int l=2*i+1,r=2*i+2,m=i;if(l<sz&&heap[l]<heap[m])m=l;if(r<sz&&heap[r]<heap[m])m=r;if(m==i)break;int t=heap[i];heap[i]=heap[m];heap[m]=t;i=m;} }
int add(int val){ push(val); while(sz>K)pop(); return heap[0]; }

Complexity

Operation	Time	Space
add	O(log k)	O(k)