Kth Largest Element in a Stream

Easyheappriority-queuedesign
Category: Heap / Priority Queue
Companies that ask this question:
AmazonFacebookGoogle
LeetCode →NeetCode →

Approach

Kth Largest Element in a Stream

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.

Implement KthLargest class:

  • KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of integers nums
  • int add(int val) Appends the integer val to the stream and returns the element representing the kth largest element in the stream

Examples

Example 1

Input:
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]

Output:
[null, 4, 5, 5, 8, 8]

Explanation:
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3);   // return 4 (stream: [2,3,4,5,8], 3rd largest = 4)
kthLargest.add(5);   // return 5 (stream: [2,3,4,5,5,8], 3rd largest = 5)
kthLargest.add(10);  // return 5 (stream: [2,3,4,5,5,8,10], 3rd largest = 5)
kthLargest.add(9);   // return 8 (stream: [2,3,4,5,5,8,9,10], 3rd largest = 8)
kthLargest.add(4);   // return 8 (stream: [2,3,4,4,5,5,8,9,10], 3rd largest = 8)

Intuition

Key Insight: We don't need to track ALL elements, only the k largest elements!

Use a min-heap of size k:

  • Heap contains the k largest elements seen so far
  • Root of heap (minimum) is the kth largest element
  • When new element arrives:
    • If heap size < k: add it
    • If val > heap root: remove root, add val
    • Otherwise: ignore val

Why this works: If we maintain exactly k largest elements, the smallest of these k elements IS the kth largest!

Pattern Recognition

This is a Min-Heap with Fixed Size problem:

  • Keep heap size = k
  • Heap property maintains ordering
  • Root gives O(1) access to answer

Approach

Min-Heap of Size k

  1. Initialize:

    • Create min-heap
    • Add all nums to heap
    • If heap size > k, remove smallest elements until size = k

  2. Add Operation:

    • If heap size < k: add val
    • Else if val > heap root:
      • Remove root (pop)
      • Add val (push)
    • Return heap root (kth largest)

  3. Time Complexity:

    • Init: O(n log k) where n = initial nums
    • Add: O(log k) for heap operations

Why not max-heap? Max-heap gives us largest element, but we need kth largest. Min-heap's root naturally gives us the "boundary" element.

Edge Cases

  • k = 1 (track maximum only)
  • Stream smaller than k initially
  • All elements same value
  • Negative numbers
  • Very large or small values

Complexity Analysis

  • Init Time: O(n log k) where n = len(nums)

    • Add n elements to heap
    • Each insertion O(log k)

  • Add Time: O(log k)

    • Pop and push operations on heap of size k

  • Space: O(k)

    • Store only k largest elements

Much better than sorting stream every time: O(n log n) per add!

Solution

java
1import java.util.PriorityQueue;
2
3class KthLargest {
4    private int k;
5    private PriorityQueue<Integer> minHeap;
6
7    public KthLargest(int k, int[] nums) {
8        this.k = k;
9        this.minHeap = new PriorityQueue<>();
10
11        // Add all elements to heap
12        for (int num : nums) {
13            minHeap.offer(num);
14        }
15
16        // Maintain heap size of k by removing smallest elements
17        while (minHeap.size() > k) {
18            minHeap.poll();
19        }
20    }
21
22    public int add(int val) {
23        // Add to heap if size < k
24        if (minHeap.size() < k) {
25            minHeap.offer(val);
26        }
27        // If val is larger than smallest in heap, replace it
28        else if (val > minHeap.peek()) {
29            minHeap.poll();
30            minHeap.offer(val);
31        }
32
33        // Root of min-heap is kth largest
34        return minHeap.peek();
35    }
36}
37
Loading visualizer...