Start learning free

LESSON 402 OF 573 · HEAP

Kth Largest Element in a Stream in Go: 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 9
    Ctrl + Enter to check

    GO · HEAP

    Kth Largest Element In A Stream in Go, explained

    Problem overview

    Numbers arrive one at a time, and after each one the program must name the k-th largest number seen so far. Equal numbers count separately: in [8, 7, 7, 5] the 2nd largest is 7, and so is the 3rd.

    The lesson asks for a KthLargest struct with a Constructor(k, nums) function, which starts the object with k and the numbers that have arrived already, and an Add(val) method, which takes in one more number and returns the current k-th largest. The statement promises that at least k numbers have arrived whenever Add returns, so the answer always exists.

    Take k = 3 and nums = [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. After Add(5) the three largest are 8, 5, 5 and the answer is 5.

    Approach: keep only the k largest

    The direct way is to store every number and sort them on each Add. That costs O(n log n) per call, where n is how many numbers have arrived, and it keeps growing.

    The better way starts from one observation: a number that is not among the k largest can never become the k-th largest again. New numbers only push the k-th largest up, never down. So everything below the top k can be thrown away for good, and the object only needs to hold k values.

    Those k values live in a min-heap, a structure whose smallest value is always at index 0 and which adds or removes a value in O(log k). The smallest of the k largest numbers is, by definition, the k-th largest, so the answer is always waiting at the top.

    When a new number arrives it is pushed in. If the heap then holds k + 1 values, the smallest of them is removed: it cannot be among the k largest. The heap never holds more than k + 1 values.

    The lesson provides the minHeap type, a slice of ints that works with Go's container/heap package. Its zero value, an empty slice, is an empty heap that is ready to use.

    Step by step

    1. Declare the field k as an int. Add needs it to know when the heap has grown too big.
    2. Declare the field largest as a minHeap, holding the k largest values with the smallest on top, at largest[0].
    3. In Constructor, create the struct with KthLargest{k: k}. The field left out, largest, starts as an empty heap.
    4. Loop over nums with range, skipping the index with _. The numbers that arrived before the object existed count like any others, and they go through the same filter.
    5. Call kthLargest.Add(num) in the loop. Reusing Add avoids repeating the push-and-trim logic. Its return value is ignored because nobody has asked yet. Go passes the address of kthLargest to the pointer method automatically, so the call changes this very struct.
    6. Return kthLargest after the loop.
    7. In Add, push val with heap.Push and the pointer &this.largest. The pointer is needed because pushing grows the slice.
    8. Trim the heap. If Len() is now above this.k, call heap.Pop(&this.largest) to remove the smallest value.
    9. Return this.largest[0], the top of the heap.

    The solution in Go

    type KthLargest struct {
    k int
    largest minHeap
}

func Constructor(k int, nums []int) KthLargest {
    kthLargest := KthLargest{k: k}
    for _, num := range nums {
        kthLargest.Add(num)
    }
    return kthLargest
}

func (this *KthLargest) Add(val int) int {
    heap.Push(&this.largest, val)
    if this.largest.Len() > this.k {
        heap.Pop(&this.largest)
    }
    return this.largest[0]
}

    A few Go details worth a closer look:

    • Pointer receivers. Add changes this.largest, so it needs the receiver *KthLargest. In Constructor, kthLargest is an addressable local variable, so kthLargest.Add(num) is shorthand for (&kthLargest).Add(num).
    • Returning the struct by value. Constructor returns a copy of kthLargest. The copy shares the heap's backing array with the local variable, which is harmless because the local variable is never used again.
    • heap.Push and heap.Pop. They take a pointer to a value that implements heap.Interface and keep the order: after either call, index 0 holds the smallest value. Reading this.largest[0] directly removes nothing.
    • The heap's own slice. A minHeap is a slice, so the heap's contents can be printed with fmt.Println. The printed order is the heap's internal order, not a sorted order. Only index 0 is guaranteed to be the smallest.

    Complexity

    Time: O(log k) per Add. The heap holds at most k + 1 values, so one push and one pop each take O(log k). Reading the top is O(1). The constructor makes one Add call for each number in nums, so it takes O(m log k) for m starting numbers.

    Space: O(k). Only the k largest values are kept, however many numbers arrive.

    Edge cases and common mistakes

    • Equal values. With k = 4 and nums = [7, 7, 7, 7, 8, 3], the calls Add(2), Add(10), Add(9), Add(9) return 7, 7, 7, 8. The heap keeps duplicates as separate values, so the second 9 pushes the last 7 out.
    • k = 1. The heap holds a single value, the largest so far. With k = 1 and no starting numbers, Add(-3), Add(-2), Add(-4), Add(0), Add(4) return -3, -2, -2, 0, 4.
    • Using a max-heap. A max-heap puts the largest value on top, which is the wrong end: the k-th largest is the smallest of the kept values.
    • Forgetting the starting numbers. They must be fed through Add too. Skipping them leaves the heap short of values.
    • Trimming with if versus for. One if is enough because each call adds exactly one value, so the heap is over by at most one.
    • Passing the heap without &. heap.Push(this.largest, val) does not compile, because the methods that grow the slice are on the pointer type.

    Testing

    Using Constructor(3, []int{4, 5, 8, 2}), the calls Add(3), Add(5), Add(10), Add(9), Add(4) return 4, 5, 5, 8, 8. Right after the constructor, the heap holds [4 5 8]: the 2 was pushed out, and the heap's top, 4, is the 3rd largest. After each of those five calls it still holds exactly 3 values.

    After 100,000 calls on an object made with k = 5, the heap holds 5 values, so the memory used does not grow with the stream.

    Keep going with Go500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free