Kth Largest Element in a Stream

NeetCode: Kth Largest Element in a Stream · LeetCode: 703. Kth Largest Element in a Stream · Difficulty: Easy

Design a class that tracks the k-th largest element of a stream. It starts with an array nums; each add(val) inserts a value and returns the current k-th largest.

Example

Text

k = 3, nums = [1, 2, 3, 3]
add(3) -> 3        add(5) -> 3        add(6) -> 3        add(7) -> 5        add(8) -> 6

Brute force: sort after every add

The class is KthLargest with LeetCode's constructor KthLargest(int k, int[] nums) and method int add(int val). Collections.reverseOrder() is a comparator that sorts in descending order.

Java

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.StringJoiner;

class KthLargest {
    private final int k;
    private final List<Integer> values = new ArrayList<>();

    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int n : nums) {
            values.add(n);
        }
    }

    public int add(int val) {
        values.add(val);
        values.sort(Collections.reverseOrder());
        return values.get(k - 1);
    }
}

public class Main {
    public static void main(String[] args) {
        KthLargest kth = new KthLargest(3, new int[]{1, 2, 3, 3});
        StringJoiner out = new StringJoiner(" ");
        for (int v : new int[]{3, 5, 6, 7, 8}) {
            out.add(String.valueOf(kth.add(v)));
        }
        System.out.println(out);
    }
}

Output

3 3 3 5 6

Time: O(n log n) per add. Space: O(n).

Optimal: a min-heap holding the k largest

Keep only the k largest values seen, in a min-heap. The smallest of them, at the top, is the k-th largest overall. When a new value arrives, add it and, if the heap has grown past k, poll the smallest. Java's PriorityQueue is a min-heap by default, so no wrapper is needed: poll removes the smallest value and peek reads it without removing it. The constructor calls add for each initial value, and a short initial array simply leaves the heap with fewer than k values until later adds fill it.

Java

import java.util.PriorityQueue;
import java.util.StringJoiner;

class KthLargest {
    private final int k;
    private final PriorityQueue<Integer> heap = new PriorityQueue<>(); // min-heap

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

    public int add(int val) {
        heap.add(val);
        if (heap.size() > k) {
            heap.poll();
        }
        return heap.peek();
    }
}

public class Main {
    public static void main(String[] args) {
        KthLargest kth = new KthLargest(3, new int[]{1, 2, 3, 3});
        StringJoiner out = new StringJoiner(" ");
        for (int v : new int[]{3, 5, 6, 7, 8}) {
            out.add(String.valueOf(kth.add(v)));
        }
        System.out.println(out);
    }
}

Output

3 3 3 5 6

Time: O(log k) per add. Space: O(k).