Start learning free

LESSON 403 OF 574 · HEAP

Kth Largest Element in a Stream in Swift: 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

    SWIFT · HEAP

    Kth Largest Element In A Stream in Swift, explained

    Problem overview

    Numbers arrive one at a time, and after each one the class must report 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 class is KthLargest. Its initializer takes k and the numbers that have already arrived, and add(_:) takes one new number and returns the current k-th largest.

    Take KthLargest(3, [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) and add(4) return 4, 5, 5, 8, 8. The last call returns 8 because 4 is smaller than the current 3rd largest and changes nothing at the top.

    The statement asks for O(log k) time per add and O(k) stored values, so the class can't keep everything and sort it again.

    Approach: keep only the k largest

    The simple way is to store every number and, on each add, sort them and read the k-th from the end. That costs O(n log n) per call, where n is how many numbers have arrived, and the storage grows without limit. With 100,000 starting numbers and 100,000 calls, that is far too slow.

    The way out is to notice which numbers can still matter. Values only ever arrive, never leave, so the k-th largest can only stay the same or go up. A number that is below the current k-th largest will never be the answer again, and it can be forgotten. Only the k largest numbers need to be kept.

    Among those k numbers, the one we need is the smallest, because the k-th largest is the smallest of the top k. That calls for a min-heap: a structure that keeps its smallest value on top, reads it in O(1), and inserts or removes in O(log size). With a heap of size k, those operations cost O(log k).

    So each add does three things. It pushes the new value. If the heap now holds k + 1 values, its smallest cannot be among the k largest, so it is removed. Then the top is the answer. In the first example the heap holds exactly three values after the initializer and after every call, and the tops the calls return are 4, 5, 5, 8, 8.

    Step by step

    1. Declare k as a let property of type Int, since add needs it later and it never changes.
    2. Declare largest as a var set to MinHeap(). The lesson provides MinHeap with insert, extract() (removes and returns the top), peek() (reads the top), count and isEmpty. It must be a var because it is a struct whose methods change it.
    3. Store the parameter with self.k = k. The parameter and the property share a name, and self.k is the property.
    4. Loop over nums: the starting numbers count like any later number, so they go through the same filter.
    5. Call add(num) and ignore its answer with _ =. This reuses the code that keeps the heap at k values, and nobody has asked a question yet.
    6. In add, insert val into largest.
    7. If largest.count is more than k, call extract(). The smallest of k + 1 values isn't among the k largest now, and since values only arrive, it won't be later.
    8. Return largest.peek(): the smallest of the k largest values is the k-th largest.

    The solution in Swift

    class KthLargest {
    let k: Int
    var largest = MinHeap()

    init(_ k: Int, _ nums: [Int]) {
        self.k = k
        for num in nums {
            _ = add(num)
        }
    }

    func add(_ val: Int) -> Int {
        largest.insert(val)
        if largest.count > k {
            largest.extract()
        }
        return largest.peek()
    }
}

    The Swift details worth a closer look:

    • Calling a method in init. The initializer calls add(num), which is a method of the object being built. Swift allows this only after every stored property has a value. Here k is assigned on the line before the loop, and largest has its default, so the call is legal. Move the loop above self.k = k and the compiler reports that self is used in the method call add before all stored properties are initialized.
    • let versus var. k is a let and can be set once, in the initializer. largest is a var because inserting changes the heap.
    • Structs and methods that change them. MinHeap is a struct, a value type. Its mutating methods work on largest in place because largest is a stored var of a class, and add can change class properties without being marked mutating.
    • _ = add(num). add returns an Int. Assigning to _ says the result is thrown away on purpose, so there is no unused-value warning.
    • The unused result of extract(). The line largest.extract() discards the removed value; only its removal matters.

    Complexity

    Time. Each add does one insert and at most one extract on a heap that never holds more than k + 1 values, so both cost O(log k). The initializer calls add once per starting number, so it takes O(m log k) for m numbers. Reading the top is O(1).

    Space. O(k). The heap holds at most k values after each call, however many numbers have arrived. Storing and sorting all the numbers would take O(n) space and O(n log n) per call.

    Edge cases and common mistakes

    • Starting with an empty array. KthLargest(1, []) works: the heap is empty, and the first add(-3) makes it [-3]. The adds -3, -2, -4, 0, 4 return -3, -2, -2, 0, 4, the running maximum.
    • Fewer than k numbers at the start. The lesson guarantees at least k numbers have arrived whenever add returns, so the heap holds exactly k by then. Without that guarantee, peek() would return the smallest value so far, which is not a k-th largest.
    • Duplicates. Equal values are separate entries. KthLargest(4, [7, 7, 7, 7, 8, 3]) followed by add(2), add(10), add(9) and add(9) returns 7, 7, 7, 8. Removing duplicates would give a wrong answer.
    • Using a max-heap. A max-heap holds the largest value on top, which is useless here. The answer sits at the bottom of the top k, so the heap must put the smallest first.
    • Keeping everything. Pushing every value and never extracting stores all the numbers and gives the wrong top.
    • A value smaller than everything in the heap. It is inserted and then immediately extracted, since it is the smallest of the k + 1. With KthLargest(3, [4, 5, 8, 2]), calling add(1) returns 4 and the heap still holds three values.

    Testing

    • KthLargest(3, [4, 5, 8, 2]) with add(3), add(5), add(10), add(9), add(4) gives 4, 5, 5, 8, 8.
    • KthLargest(1, []) with add(-3), add(-2), add(-4), add(0), add(4) gives -3, -2, -2, 0, 4.
    • KthLargest(4, [7, 7, 7, 7, 8, 3]) with add(2), add(10), add(9), add(9) gives 7, 7, 7, 8.
    • KthLargest(2, [8, 7, 7, 5]) with add(1) gives 7.
    Keep going with Swift500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free