Kth Largest Element in a Stream

Difficulty: Easy | Acceptance: 61.00% | Paid: No Topics: Tree, Design, Binary Search Tree, Heap (Priority Queue), Binary Tree, Data Stream

You are part of a university admissions office and are tasked with simplifying the admissions process. You have a list of the scores of all the students who have applied, and you need to select the top k students. However, new students are constantly applying, and you need to be able to quickly determine the kth largest score at any point in time.

Design a class to find the kth largest element in a stream. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Implement KthLargest class:

KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of integers nums. int add(int val) Appends the integer val to the stream and returns the element representing the kth largest element in the stream.

• Examples
• Constraints
• Approach 1: Sorting (Brute Force)
• Approach 2: Min-Heap (Priority Queue)
• Approach 3: Binary Search Tree (Order Statistic Tree)

Examples

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

1 <= k <= 10⁴
0 <= nums.length <= 10⁴
-10⁴ <= nums[i] <= 10⁴
-10⁴ <= val <= 10⁴
At most 10⁴ calls will be made to add.
It is guaranteed that there will be at least k elements in the array when you search for the kth element.

Approach 1: Sorting (Brute Force)

Intuition The most straightforward way to find the kth largest element is to maintain a sorted list of all numbers. When a new number is added, we insert it into the list, sort the list again, and then return the element at the index length - k.

Steps

• Initialize the class with k and the initial nums array. Sort the nums array immediately.
• In the add method:
  • Append the new val to the list.
  • Sort the list.
  • Return the element at index list.length - k.

class KthLargest:
    def __init__(self, k: int, nums: list[int]):
        self.k = k
        self.nums = nums
        self.nums.sort()

    def add(self, val: int) -> int:
        self.nums.append(val)
        self.nums.sort()
        return self.nums[-self.k]

Complexity

• Time: O(N log N) for the add operation, where N is the total number of elements. Sorting takes O(N log N).
• Space: O(N) to store the stream of numbers.
• Notes: This approach is simple but inefficient for large streams or frequent add calls due to the repeated sorting cost.

Approach 2: Min-Heap (Priority Queue)

Intuition We only need to keep track of the top k largest elements at any time. If we maintain a min-heap of size k, the smallest element in this heap (the root) will be the kth largest element in the stream. Any element smaller than the root can be ignored because it will never be in the top k.

Steps

• Initialize a min-heap.
• Add all numbers from nums to the heap.
• If the heap size exceeds k, remove the smallest element (root) to maintain size k.
• In the add method:
  • Add the new val to the heap.
  • If the heap size exceeds k, remove the root.
  • Return the root element.

import heapq

class KthLargest:
    def __init__(self, k: int, nums: list[int]):
        self.k = k
        self.heap = nums
        heapq.heapify(self.heap)
        while len(self.heap) > k:
            heapq.heappop(self.heap)

    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)
        if len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

Complexity

• Time: O(log k) for the add operation. Insertion and extraction in a heap of size k takes O(log k).
• Space: O(k) to store the heap.
• Notes: This is the optimal approach for this problem, significantly faster than sorting when k is much smaller than the total number of elements.

Approach 3: Binary Search Tree (Order Statistic Tree)

Intuition We can use a Binary Search Tree (BST) to store the elements. By augmenting the BST nodes to store the size of the subtree rooted at that node, we can efficiently find the kth largest element. The kth largest element corresponds to finding the node at a specific rank when traversing the tree in reverse in-order (right-root-left).

Steps

• Define a TreeNode class with val, left, right, and size (count of nodes in subtree).
• Implement insert to add a value to the BST and update the size of ancestors.
• Implement findKth to traverse the tree. If the right subtree has R nodes:
  • If k == R + 1, current node is the answer.
  • If k <= R, search in the right subtree.
  • If k > R + 1, search in the left subtree with updated rank k - (R + 1).

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
        self.size = 1

class KthLargest:
    def __init__(self, k: int, nums: list[int]):
        self.k = k
        self.root = None
        for num in nums:
            self.root = self._insert(self.root, num)

    def _insert(self, node: TreeNode, val: int) -> TreeNode:
        if not node:
            return TreeNode(val)
        if val <= node.val:
            node.left = self._insert(node.left, val)
        else:
            node.right = self._insert(node.right, val)
        node.size += 1
        return node

    def _find_kth(self, node: TreeNode, k: int) -> int:
        if not node:
            return 0
        right_size = node.right.size if node.right else 0
        if k == right_size + 1:
            return node.val
        elif k <= right_size:
            return self._find_kth(node.right, k)
        else:
            return self._find_kth(node.left, k - right_size - 1)

    def add(self, val: int) -> int:
        self.root = self._insert(self.root, val)
        return self._find_kth(self.root, self.k)

Complexity

• Time: O(H) for both add and finding the element, where H is the height of the tree. In the worst case (unbalanced tree), this is O(N). On average (balanced tree), this is O(log N).
• Space: O(N) to store the tree nodes.
• Notes: This approach is more complex to implement than the Heap approach and offers similar average-case performance but worse worst-case performance unless the tree is self-balancing (e.g., AVL, Red-Black).