Start learning free

LESSON 398 OF 569 · HEAP

Kth Largest Element in a Stream in Java: Min-Heap of Size k

Practice O(log k) Read the full explanation

Numbers arrive one at a time, like the scores of a game coming in live, and after each one you must say which number is in k-th place: the k-th largest of all the numbers so far. Equal numbers count separately: in [8, 7, 7, 5] the 2nd largest is 7, and so is the 3rd.

Implement the class KthLargest:

  • KthLargest(k, nums) creates the object with the integer k and the numbers nums that have arrived already.
  • add(val) adds the number val and returns the k-th largest of all the numbers so far.

The input describes one new KthLargest and a list of add calls, applied in order. The output lists what the add calls return, in order. The starter code spells these names the way your language does.

 

Example 1:

Input:
    KthLargest(3, [4, 5, 8, 2])
    add(3)
    add(5)
    add(10)
    add(9)
    add(4)
Output: 4 5 5 8 8
Explanation:
After add(3) the numbers are [2, 3, 4, 5, 8]: the three largest are 8, 5 and 4, so the 3rd largest is 4.
After add(5): 8, 5, 5 are the three largest, so the answer is 5.
After add(10): 10, 8, 5, answer 5. After add(9): 10, 9, 8, answer 8.
add(4) changes nothing at the top: 4 is smaller than the 3rd largest, so the answer stays 8.

Example 2:

Input:
    KthLargest(1, [])
    add(-3)
    add(-2)
    add(-4)
    add(0)
    add(4)
Output: -3 -2 -2 0 4
Explanation: With k = 1, add returns the largest number so far. No numbers have arrived when the object is created, so after the first add, -3 is the only one, and the largest.

Example 3:

Input:
    KthLargest(4, [7, 7, 7, 7, 8, 3])
    add(2)
    add(10)
    add(9)
    add(9)
Output: 7 7 7 8
Explanation: Equal numbers count separately: in [8, 7, 7, 7, 7, 3, 2] the 4th largest is the third 7. After 10 and 9 arrive the four largest are 10, 9, 8, 7. The second 9 pushes the last 7 out: 10, 9, 9, 8, so the answer becomes 8.

 

Constraints:

  • 1 <= k <= 105
  • 0 <= nums.length <= 105
  • -104 <= nums[i], val <= 104
  • At most 105 calls are made to add.
  • When add returns, at least k numbers have arrived, so the k-th largest always exists.

Note:

Each add should run in O(log k) time, and the object should keep only O(k) values. Sorting all the numbers again on every add would take O(n log n) time per call, where n is how many numbers have arrived so far.

STEP 1 OF 8
    Ctrl + Enter to check

    JAVA · HEAP

    Kth Largest Element In A Stream in Java, explained

    Problem overview

    Numbers arrive one at a time, and after each one you must report the k-th largest of everything seen so far. The class KthLargest has a constructor KthLargest(int k, int[] nums), which receives the numbers that arrived before the first question, and a method add(int val), which takes a new number and returns the current k-th largest. Equal numbers count separately: among 8, 7, 7, 5 the 2nd largest is 7, and so is the 3rd.

    Take k = 3 and nums = [4, 5, 8, 2]. The five calls add(3), add(5), add(10), add(9) and add(4) return the values below. After add(3) the numbers are 2, 3, 4, 5, 8, whose three largest are 8, 5 and 4, so the answer is 4.

    4 5 5 8 8

    The constraints guarantee that at least k numbers have arrived whenever add returns, so the k-th largest always exists.

    Approach: keep only the k largest

    The simple approach stores every number in a list and sorts it on each add, then reads the element at position k from the top. Sorting costs O(n log n) per call, where n is how many numbers have arrived, and with up to 100,000 numbers and 100,000 calls that is far too slow.

    The key observation is that most numbers can be forgotten. Values only ever arrive; none leave. So a number that is not among the k largest now can never become the k-th largest later, because the k numbers above it stay above it. Only the k largest values matter.

    Which structure keeps the k largest values and tells you the smallest of them quickly? A min-heap. Its top is the smallest value it holds, and since it holds the k largest values, that smallest one is exactly the k-th largest. Java's PriorityQueue is a min-heap when you give it no comparator, so new PriorityQueue<>() is what the lesson needs.

    Each add then does two things. It pushes the new value, and if the heap now holds k + 1 values, it removes the smallest. The removed value cannot be among the k largest, now or later. The heap never holds more than k + 1 values, so a push and a pop each cost O(log k), and reading the top is O(1).

    Step by step

    1. Declare the field k. add needs it long after the constructor has finished, so it is a field, not just a parameter.
    2. Declare the heap largest as a PriorityQueue<Integer> that starts empty.
    3. Store the parameter in the field with this.k = k. The parameter and the field share a name, and this.k names the field.
    4. Loop over nums. The starting numbers count like any later one, so they must pass the same filter: kept only while they are among the k largest.
    5. Call add(num) in the loop. It already keeps the heap at size k, so the constructor reuses it instead of repeating that logic. The returned value is ignored, because nobody has asked a question yet.
    6. In add, put val in the heap.
    7. If the heap holds more than k values, call poll(). It removes the smallest, which cannot be one of the k largest.
    8. Return peek(). It reads the top without removing it. With exactly k values in the heap, the top is the k-th largest.

    The solution in Java

    class KthLargest {
    int k;
    PriorityQueue<Integer> largest = new PriorityQueue<>();

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

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

    A few details about the Java code:

    • Boxing. A PriorityQueue<Integer> holds Integer objects, not int. largest.add(val) boxes the int, and return largest.peek() unboxes it back to the int that add promises.
    • add versus poll and peek. add inserts, poll() removes and returns the smallest, and peek() returns the smallest without removing it. On an empty heap, poll() and peek() return null, and unboxing that null into an int would throw a NullPointerException. The constraints rule that case out.
    • Calling a method from the constructor. The constructor calls add, a method of the same class. That works because the fields k and largest are set before the loop runs: largest is initialised in its declaration, and this.k = k comes first in the constructor.
    • The heap's order is not sorted order. Only the top is guaranteed to be the smallest. Printing a PriorityQueue or iterating over it shows the internal array, not a sorted list. Use only peek() and poll() to read it.

    Seeing it work

    Create the object and call add the way the lesson's tests do. After the constructor has processed [4, 5, 8, 2] with k = 3, the heap holds three values: 2 was pushed out, and the top is 4. Each add then returns the heap's top.

    KthLargest kl = new KthLargest(3, new int[]{4, 5, 8, 2});
for (int v : new int[]{3, 5, 10, 9, 4}) {
    System.out.print(kl.add(v) + " ");
}

    4 5 5 8 8

    With k = 1 the heap holds a single value, the largest so far. Starting from an empty array, the five adds of -3, -2, -4, 0, 4 give:

    -3 -2 -2 0 4

    Duplicates matter. With k = 4 and starting numbers [7, 7, 7, 7, 8, 3], the adds of 2, 10, 9, 9 give:

    7 7 7 8

    The first add returns 7, because the four largest of 8, 7, 7, 7, 7, 3, 2 are 8 and three 7s. After 10 and 9 arrive the four largest are 10, 9, 8, 7. The second 9 pushes the last 7 out, leaving 10, 9, 9, 8, so the answer becomes 8.

    Complexity

    Time per add: O(log k). The heap holds at most k + 1 values. A push and a pop each move a value up or down a tree of height about log k, and peek() is constant time. The constructor makes one add per starting number, so it takes O(m log k) for m starting numbers.

    Space: O(k). The object keeps only the k largest values, however many numbers have arrived. Re-sorting everything, by comparison, takes O(n log n) time per call and O(n) space.

    Edge cases and common mistakes

    • Using a max-heap. A max-heap would put the largest value on top, which is not the answer. The k-th largest is the smallest of the k largest, so the min-heap is the right order. In Java the default PriorityQueue is already a min-heap, so no comparator is needed.
    • Removing the wrong end. After the push, the value to drop is the smallest of k + 1. Dropping the largest would throw away exactly the values that matter.
    • Forgetting the starting numbers. Skipping the loop in the constructor leaves the heap empty, and the first add then answers from too few values.
    • Using >= instead of >. The condition is size() > k. With >= the heap would hold only k - 1 values and its top would be the (k-1)-th largest.
    • Treating equal numbers as one. Each copy takes its own place in the heap. In the third run above, the three 7s each count.
    • k = 1. The heap holds one value, and every add keeps the larger of the old value and the new one.
    Keep going with Java500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free