703. Kth Largest Element in a Stream

703. Kth Largest Element in a Stream#

Difficulty: Easy

Link to Problem: To see the Kth Largest Element in a Stream problem on LeetCode, click here!

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 the KthLargest class with the following methods:

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.

Constraints:

1 <= k <= \(10^4\)
0 <= nums.length <= \(10^4\)
\(-10^4\) <= nums[i] <= \(10^4\)
\(-10^4\) <= val <= \(10^4\)
At most \(10^4\) 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.

Probelm Explanation:#

The problem you’re trying to solve is to design a class called KthLargest that can efficiently find the kth largest element in a stream of integers. You need to support two main operations:

To solve this problem efficiently, you can use a min-heap data structure. Here’s how the approach works:

Initialize a min-heap to store the k largest elements. Initially, it’s empty.
In the __init__ method, populate the min-heap with the first k elements from nums. This ensures that you have the k largest elements seen so far.
Whenever you add a new element to the stream using the add method, follow these steps:
- Add the new element to the min-heap.
- If the size of the min-heap exceeds k, remove the smallest element from the min-heap. This ensures that you always have the k largest elements in the min-heap.
- The smallest element in the min-heap (the root) will always represent the kth largest element in the stream.

Here’s why this approach works efficiently:

By using a min-heap, you can quickly maintain the k largest elements, and finding the smallest element in the heap (the root) takes constant time.
When you add a new element, the min-heap’s size is kept at most k, which ensures that you only track the k largest elements and discard the smaller ones.
The time complexity for adding an element is O(log k), which is very efficient compared to sorting the entire stream, which would be O(n log n).
This approach meets the constraints of the problem, including handling large streams and having a low time complexity for both initialization and adding elements.

In summary, the min-heap approach efficiently tracks the kth largest element in the stream by maintaining a heap of the k largest elements seen so far, updating it as new elements are added. This approach provides a fast and scalable solution to the problem.

Solution:#

Here’s a Python function to implement this algorithm:

import heapq
from typing import List  # Import the List type from the typing module

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        # Initialize the KthLargest object with k and nums.
        self.k = k
        # Create a min-heap to store the k largest elements.
        self.min_heap = []
        # Populate the min-heap with the first k elements from nums.
        for num in nums:
            self.add(num)

    def add(self, val: int) -> int:
        # Add val to the min-heap.
        heapq.heappush(self.min_heap, val)
        # If the size of the min-heap exceeds k, remove the smallest element.
        if len(self.min_heap) > self.k:
            heapq.heappop(self.min_heap)
        # The root of the min-heap is the kth largest element.
        return self.min_heap[0]

Explanation:#

We start by importing the necessary modules:
- heapq: This module provides functions to create and manipulate heaps.
- List from the typing module: This is used to specify the type of the nums parameter.
We define the KthLargest class:
- The __init__ method initializes the object with the integer k and the list of integers nums.
- It also creates an empty min-heap called self.min_heap to store the k largest elements.
- It populates the min-heap with the first k elements from nums by calling the add method.
The add method:
- Adds the new integer val to the min-heap using heapq.heappush. This maintains the min-heap property.
- Checks if the size of the min-heap exceeds k. If it does, it removes the smallest element (the k+1th largest) using heapq.heappop.
- Finally, it returns the smallest element in the min-heap, which is always the kth largest element.

Overall, this code implements a class that efficiently finds the kth largest element in a stream of integers by maintaining a min-heap of the k largest elements seen so far.

Test cases:#

Here’s how you can use this solution:

# Example 1: 

# Example usage:
kthLargest = KthLargest(3, [4, 5, 8, 2])
print(kthLargest.add(3))  # Output: 4
print(kthLargest.add(5))  # Output: 5
print(kthLargest.add(10)) # Output: 5
print(kthLargest.add(9))  # Output: 8
print(kthLargest.add(4))  # Output: 8

Time and Space Complexity Analysis#

Let’s analyze the time and space complexity of the KthLargest class using the min-heap approach:

Time Complexity:

__init__ method:
- In the __init__ method, you iterate over the first k elements in nums and add them to the min-heap. Each heapq.heappush operation takes O(log k) time.
- Therefore, the time complexity of the __init__ method is O(k * log k).
add method:
- In the add method, you perform the following operations:
  - heapq.heappush: O(log k) to add the new element to the min-heap.
  - If the size of the min-heap exceeds k, you perform heapq.heappop, which is also O(log k).
  - Finally, you return the smallest element from the min-heap, which is O(1) because it’s always the root.
- Overall, the time complexity of the add method is O(log k).
Overall, if you make n calls to the add method, the total time complexity is O(n * log k), where n is the total number of elements added to the stream.

Space Complexity:

The space complexity is determined by the space used to store the min-heap and the instance variables.
The min-heap stores at most k elements at any given time, so its space complexity is O(k).
The instance variables (such as self.k and self.min_heap) have constant space requirements.
Therefore, the overall space complexity of the KthLargest class is O(k).

In summary, the time complexity of the KthLargest class is O(n * log k) for n add operations, and the space complexity is O(k), where k is the parameter passed during initialization. This implementation efficiently maintains the kth largest element in the stream while meeting the problem’s constraints.

Challenging Exercises:#

Dynamic k: Modify the KthLargest class to support dynamic changes in the value of k. Implement a method update_k that allows changing the value of k during the lifetime of the object. Ensure that the object can still correctly find the kth largest element based on the updated value of k.
Implement kth Smallest: Create a new class called KthSmallest that finds the kth smallest element in a stream instead of the kth largest. You may need to modify the data structure used in the implementation.