Kth Largest Element in a Stream

Problem

https://leetcode.com/problems/kth-largest-element-in-a-stream/

You are part of a university admissions office and need to keep track of the kth highest test score from applicants in real-time. This helps to determine cut-off marks for interviews and admissions dynamically as new applicants submit their scores.

You are tasked to implement a class which, for a given integer k, maintains a stream of test scores and continuously returns the kth highest test score after a new score has been submitted. More specifically, we are looking for the kth highest score in the sorted list of all scores.

Implement the KthLargest class:

• KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of test scores nums.

• int add(int val) Adds a new test score val to the stream and returns the element representing the k:sup:`th` largest element in the pool of test scores so far.

Example 1:

Input:

[“KthLargest”, “add”, “add”, “add”, “add”, “add”] [[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]

Output: [null, 4, 5, 5, 8, 8]

Explanation:

KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);

kthLargest.add(3); // return 4

kthLargest.add(5); // return 5

kthLargest.add(10); // return 5

kthLargest.add(9); // return 8

kthLargest.add(4); // return 8

Example 2:

Input:

[“KthLargest”, “add”, “add”, “add”, “add”] [[4, [7, 7, 7, 7, 8, 3]], [2], [10], [9], [9]]

Output: [null, 7, 7, 7, 8]

Explanation:

KthLargest kthLargest = new KthLargest(4, [7, 7, 7, 7, 8, 3]); kthLargest.add(2); // return 7 kthLargest.add(10); // return 7 kthLargest.add(9); // return 7 kthLargest.add(9); // return 8

Constraints:

• 0 <= nums.length <= 10:sup:`4`

• 1 <= k <= nums.length + 1

• -10:sup:`4`<= nums[i] <= 10:sup:`4`

• -10:sup:`4`<= val <= 10:sup:`4`

• At most 10:sup:`4` calls will be made to add.

Pattern

Tree, Design, Binary Search Tree, Heap (Priority Queue), Binary Tree, Data Stream

Approaches

Min Heap

Explanation

The key idea is to maintain a min heap of size \(k\). The top of the heap will always be the \(k`th largest element in the heap. When adding a new element, if the heap size is less than :math:`k\), we simply add the new element. If the heap size is \(k\) add the new element first, then pop the top of the heap—the \(k + 1\) largest element.

Code

import heapq


class KthLargest:
    """Tracks the kth largest element in a stream of values."""

    def __init__(self, k: int, nums: list[int]):
        self.k = k
        self.heap: list[int] = []
        for n in nums:
            self.add(n)

    def add(self, val: int) -> int:
        """Add ``val`` and return the kth largest element."""
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

Test

>>> from kth_largest_element_in_a_stream__min_heap import KthLargest
>>> kl = KthLargest(3, [4, 5, 8, 2])
>>> kl.add(3)
4
>>> kl.add(5)
5
>>> kl.add(10)
5
>>> kl.add(9)
8
>>> kl.add(4)
8

Complexity

\(n\) is the number of element in the stream and \(k\) is the size of the heap

Measure	Complexity	Notes
Time	\(O(n log k)\)	\(n\) additions, each add takes \(O(log k)\) time to maintain the heap
Auxiliary Space	\(O(k)\)	min heap of size \(k\)

class kth_largest_element_in_a_stream__min_heap.KthLargest(k: int, nums: list[int])

Bases: object

Tracks the kth largest element in a stream of values.

add(val: int) → int

Add val and return the kth largest element.