Start learning free

LESSON 397 OF 568 · HEAP

Kth Largest Element in a Stream in C++: O(log k) with a Min-Heap

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

    C++ · HEAP

    Kth Largest Element In A Stream in C++, explained

    Problem overview

    Numbers arrive one at a time, and after each one the class KthLargest must report the k-th largest of everything seen so far. Equal numbers count separately: among 8, 7, 7, 5 the 2nd largest is 7, and so is the 3rd. The constructor receives k and the numbers that have already arrived, and add(val) takes a new number and returns the current k-th largest. The constraints promise that at least k numbers have arrived whenever add returns, so the answer always exists.

    In the first example, k = 3 and the starting numbers are [4, 5, 8, 2]. After add(3) the numbers are 2, 3, 4, 5, 8; the three largest are 8, 5 and 4, so the answer is 4. The five calls add(3), add(5), add(10), add(9), add(4) return 4 5 5 8 8.

    The statement asks for O(log k) time per add and O(k) memory.

    Approach: keep only the k largest, in a min-heap

    The simple solution stores every number and sorts them again after each add. That costs O(n log n) per call, where n is how many numbers have arrived, and it keeps all of them in memory.

    The better idea starts from an observation: a number that is not among the k largest right now can never become the answer. Numbers only arrive, so the k-th largest can only stay the same or go up, and a smaller number only falls further behind. So the class needs to keep just k numbers.

    Which structure keeps "the k largest" and gives the smallest of them at once? A min-heap. The smallest of the k largest values is exactly the k-th largest, and a min-heap has its smallest value on top. When a new number arrives, push it. If the heap now holds k + 1 values, the smallest of them cannot be among the k largest, so remove the top. Both the push and the pop cost O(log k), because the heap never holds more than k + 1 values.

    Here is the heap for k = 3 after the constructor receives [4, 5, 8, 2], and then after the first three calls (the values kept, smallest first):

    after the constructor: 4 5 8
after add(3):          4 5 8
after add(5):          5 5 8
after add(10):         5 8 10

    The 2 never stays in the heap, and 3 is pushed and immediately removed again because it is the smallest of four. The top of the heap, the first number in each row, is the answer: 4, 5, 5.

    Step by step

    1. Declare the field k. add needs it to know when the heap has grown too big.
    2. Declare largest as a priority_queue<int, vector<int>, greater<int>>. By default a priority_queue is a max-heap. The third template argument, greater<int>, reverses the order so that the smallest value is on top.
    3. In the constructor, store the parameter with this->k = k. The parameter and the field have the same name, and this-> picks the field.
    4. Loop over nums and call add(num) for each value. The starting numbers go through the same filter as every later one, so the constructor does not repeat the "keep only k" logic. The returned value is ignored; nobody has asked yet.
    5. In add, push val into the heap.
    6. If the heap holds more than k values, pop. The smallest of k + 1 values is out of the top k for good.
    7. Return largest.top(). The heap holds exactly k values at this point, and its smallest is the k-th largest.

    The solution in C++

    class KthLargest {
public:
    int k;
    priority_queue<int, vector<int>, greater<int>> largest;

    KthLargest(int k, vector<int>& nums) {
        this->k = k;
        for (int num : nums) {
            add(num);
        }
    }

    int add(int val) {
        largest.push(val);
        if (largest.size() > k) {
            largest.pop();
        }
        return largest.top();
    }
};

    The C++ details worth a closer look:

    • The three template arguments. priority_queue<T, Container, Compare> stores values of type T in a Container (a vector here) and orders them with Compare. The default comparison puts the largest value on top. Passing greater<int> puts the smallest on top, and to name the comparison you must also name the container in the middle.
    • pop() returns nothing. It only removes the top. That is fine here, because the removed value is thrown away. top() reads without removing.
    • top() on an empty heap is undefined behaviour. add always pushes before it calls top(), so the heap is never empty at that point, including when the constructor calls add with fewer than k numbers.
    • Signed against unsigned. largest.size() is unsigned, and k is an int. The comparison works because k is at least 1, so converting it to unsigned changes nothing.
    • The constructor takes vector<int>&. The vector is passed by reference, so a list of up to 100,000 numbers is not copied.

    A run on the first example (k = 3, starting numbers [4, 5, 8, 2]) prints 4 5 5 8 8 for the five calls. After the constructor, the heap holds 3 values and its top is 4.

    Complexity

    Time: O(log k) per add. One push and at most one pop on a heap of at most k + 1 values, each O(log k), plus a constant-time top(). The constructor makes one add call per starting number, so building the object with m numbers costs O(m log k).

    Space: O(k). The heap never holds more than k + 1 values, however many numbers have arrived. Sorting everything keeps all n.

    Edge cases and common mistakes

    • Starting with an empty list. With k = 1 and no numbers, the first add(-3) returns -3, and the calls add(-3), add(-2), add(-4), add(0), add(4) return -3 -2 -2 0 4.
    • Duplicates. With k = 4 and starting numbers [7, 7, 7, 7, 8, 3], the heap holds 7, 7, 7, 8 and the calls add(2), add(10), add(9), add(9) return 7 7 7 8. Do not remove duplicates: equal values count separately.
    • A max-heap by mistake. With the default priority_queue<int>, the top is the largest value, not the k-th. The greater<int> argument is what makes the top the one to compare against.
    • Forgetting the filter in the constructor. Pushing the starting numbers without popping would leave the heap with more than k values, and its top would be too small.
    • Popping when the size equals k. The condition is "more than k". Popping at exactly k would drop a value that belongs to the top k.
    • Returning the largest. The answer is the smallest of the kept values, the top of a min-heap. With k = 3, after add(10) the heap holds 5, 8, 10 and the answer is 5, not 10.

    Testing

    • k = 3, numbers [4, 5, 8, 2], then add 3, 5, 10, 9, 4 gives 4 5 5 8 8.
    • k = 1, numbers [], then add -3, -2, -4, 0, 4 gives -3 -2 -2 0 4.
    • k = 4, numbers [7, 7, 7, 7, 8, 3], then add 2, 10, 9, 9 gives 7 7 7 8.
    • k = 1, numbers [5], then add 5, 1, 6 gives 5 5 6.
    Keep going with C++500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free