Problem Description

You need to design a class that tracks the kth highest test score from a stream of scores in real-time. This is useful for a university admissions office to dynamically determine cut-off marks as new applicants submit their scores.

The task is to implement a KthLargest class with the following functionality:

1. Constructor KthLargest(int k, int[] nums):

   • Initializes the object with an integer k (the position of the highest score you want to track)
   • Takes an initial array of test scores nums

2. Method int add(int val):

   • Adds a new test score val to the existing collection of scores
   • Returns the kth largest element after adding this new score

Important Note: The kth largest element means the kth highest score when all scores are sorted in descending order. For example, if scores are [100, 90, 85, 80, 75] and k=3, the 3rd largest element is 85.

The solution uses a min heap (priority queue) of size k to efficiently maintain the k largest elements. The root of this min heap will always be the kth largest element:

• During initialization, elements from nums are added to the heap
• When adding a new score, it's pushed to the heap
• If the heap size exceeds k, the smallest element (root of min heap) is removed
• The root of the heap (accessible via self.min_q[0]) is always the kth largest element

This approach ensures that we only keep track of the k largest elements at any time, making it memory efficient and providing O(log k) time complexity for each add operation.

Intuition

The key insight is that we don't need to keep track of all scores - we only care about the k largest ones. If we have more than k scores, any score smaller than the kth largest will never be relevant to our answer.

Think of it like maintaining a "top k leaderboard". If k=3 and we want to track the top 3 scores, once we have scores like [100, 90, 85], any new score below 85 won't affect who's in 3rd place unless it's higher than 85.

A min heap is perfect for this because:

1. We can limit its size to exactly k elements
2. The smallest element in our "top k" collection (which is the kth largest overall) sits at the root of the min heap
3. When a new score comes in, we add it to the heap and if we now have more than k elements, we remove the smallest one

Why specifically a min heap and not a max heap? Because we want quick access to the kth largest element (the "cutoff" score). In a min heap of size k containing the k largest elements:

• The root is the minimum among these k elements
• This minimum is exactly the kth largest element overall
• We can access it in O(1) time with min_q[0]

For example, if k=3 and we have scores [100, 95, 90, 85, 80], our min heap would contain [90, 95, 100]. The root (90) is the 3rd largest score. When we add a new score like 92, the heap becomes [90, 92, 95, 100], we pop 90, and now [92, 95, 100] gives us the new 3rd largest score of 92.

This approach elegantly handles the dynamic nature of the problem - as new scores arrive, we automatically maintain exactly what we need without sorting the entire list each time.

Learn more about Tree, Binary Search Tree, Binary Tree, Data Stream and Heap (Priority Queue) patterns.

Solution Approach

The solution uses a min heap (priority queue) to efficiently maintain the k largest elements from the stream of scores.

Data Structure Choice

We use Python's heapq module which implements a min heap. The key property we leverage is that in a min heap of size k containing the k largest elements, the root element is always the kth largest overall.

Implementation Details

1. Initialization (__init__ method):

• Store the value of k as an instance variable
• Initialize an empty min heap self.min_q = []
• Process the initial array nums by calling self.add(x) for each element
• This ensures the heap is properly initialized with at most k elements

2. Adding Elements (add method): The add method performs three key operations:

heappush(self.min_q, val)  # Step 1: Add new element
if len(self.min_q) > self.k:  # Step 2: Check size
    heappop(self.min_q)  # Step 3: Remove smallest if needed
return self.min_q[0]  # Step 4: Return kth largest

• Step 1: Push the new value onto the heap using heappush. This maintains the heap property in O(log n) time
• Step 2: Check if the heap size exceeds k
• Step 3: If size > k, remove the smallest element (root) using heappop. This ensures we keep only the k largest elements
• Step 4: Return the root element self.min_q[0], which is the kth largest element

Time Complexity Analysis

• Initialization: O(n × log k) where n is the length of the initial nums array. Each element is added using add, which takes O(log k) time
• Add operation: O(log k) for both heappush and potentially heappop operations
• Getting kth largest: O(1) as we simply access the root element

Space Complexity

O(k) - we maintain at most k elements in the heap at any time, regardless of how many total elements are processed.

Example Walkthrough

If k=3 and initial nums=[4, 5, 8, 2]:

1. After initialization, heap contains [4, 5, 8] (the 3 largest)
2. Call add(3): heap becomes [3, 4, 5, 8], pop smallest → [4, 5, 8], return 4
3. Call add(10): heap becomes [4, 5, 8, 10], pop smallest → [5, 8, 10], return 5

The beauty of this approach is that it automatically maintains exactly what we need - the k largest elements - with optimal time and space complexity.

Example Walkthrough

Let's walk through a concrete example with k=3 (tracking the 3rd highest score).

Initial Setup:

• k = 3
• nums = [4, 5, 8, 2]

Step 1: Initialization We process each element from nums:

• Add 4: heap = [4], size = 1 ≤ k, no removal needed
• Add 5: heap = [4, 5], size = 2 ≤ k, no removal needed
• Add 8: heap = [4, 5, 8], size = 3 = k, no removal needed
• Add 2: heap = [2, 4, 5, 8], size = 4 > k, remove smallest (2) → heap = [4, 5, 8]

After initialization, our min heap contains [4, 5, 8] where 4 is at the root.

Step 2: Add new score 3

• Push 3 to heap: [3, 4, 5, 8]
• Size = 4 > k = 3, so pop smallest: remove 3 → heap = [4, 5, 8]
• Return heap[0] = 4 (the 3rd largest among all scores: 8, 5, 4, 3, 2)

Step 3: Add new score 10

• Push 10 to heap: [4, 5, 8, 10]
• Size = 4 > k = 3, so pop smallest: remove 4 → heap = [5, 8, 10]
• Return heap[0] = 5 (the 3rd largest among all scores: 10, 8, 5, 4, 3, 2)

Step 4: Add new score 9

• Push 9 to heap: [5, 8, 9, 10]
• Size = 4 > k = 3, so pop smallest: remove 5 → heap = [8, 9, 10]
• Return heap[0] = 8 (the 3rd largest among all scores: 10, 9, 8, 5, 4, 3, 2)

Key Observation: The min heap always maintains exactly the k largest elements, and its root (smallest element in the heap) is precisely the kth largest element overall. This gives us O(1) access to the answer we need.

Solution Implementation

Time and Space Complexity

Time Complexity:

• Initialization (__init__): O(n * log k) where n is the length of the initial nums array. Each element in nums calls the add method, which performs a heap push operation O(log k) and potentially a heap pop operation O(log k).
• Add operation (add): O(log k) for each call. The heappush operation takes O(log k) time since the heap size is at most k+1, and the heappop operation (when needed) also takes O(log k) time.

Space Complexity: O(k)

• The min heap self.min_q stores at most k elements at any time. After each add operation, if the heap size exceeds k, the smallest element is removed, maintaining the heap size at exactly k elements (assuming at least k elements have been added).
• The additional space used for the integer k and other variables is O(1), which doesn't affect the overall space complexity.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Not Handling Edge Case When Initial Array Has Fewer Than k Elements

The Pitfall: The most common mistake is assuming the heap will always have exactly k elements. When the initial nums array has fewer than k elements, or even when it's empty, accessing self.min_heap[0] in the add method will raise an IndexError.

Example of the Problem:

# This will crash!
kthLargest = KthLargest(3, [4, 5])  # Only 2 elements, but k=3
result = kthLargest.add(3)  # IndexError when accessing self.min_heap[0]

The Solution: Check if the heap has at least k elements before returning the root:

def add(self, val: int) -> int:
    heapq.heappush(self.min_heap, val)
  
    if len(self.min_heap) > self.k:
        heapq.heappop(self.min_heap)
  
    # Only return if we have at least k elements
    if len(self.min_heap) == self.k:
        return self.min_heap[0]
    else:
        # Could return None or raise an exception
        # Based on problem requirements
        return None  # or return -1 as a sentinel value

2. Using a Max Heap Instead of Min Heap

The Pitfall: Developers might intuitively think "I need the k largest, so I should use a max heap." This leads to incorrect logic where they try to negate values or use complex workarounds.

Example of the Problem:

# Wrong approach - negating values for max heap
def add(self, val: int) -> int:
    heapq.heappush(self.min_heap, -val)  # Negating for max heap
    if len(self.min_heap) > self.k:
        heapq.heappop(self.min_heap)
    return -self.min_heap[0]  # Need to negate back

Why This is Wrong: While this might seem to work, it complicates the logic unnecessarily and can lead to bugs with edge cases (like integer overflow with very large negative numbers).

The Correct Understanding: A min heap of size k containing the k largest elements naturally has the kth largest at its root. This is because all other elements in the heap are larger than or equal to the root.

3. Forgetting to Handle Duplicate Values

The Pitfall: Some implementations try to avoid duplicates or handle them specially, which is unnecessary and incorrect for this problem.

Example of Incorrect Thinking:

def add(self, val: int) -> int:
    # Wrong: trying to avoid duplicates
    if val not in self.min_heap:
        heapq.heappush(self.min_heap, val)
    # This changes the problem semantics!

The Solution: Simply treat duplicates as separate values. If you have scores [90, 90, 85, 80] and k=2, the 2nd largest is 90 (not 85).

4. Inefficient Initialization

The Pitfall: Some developers try to optimize initialization by sorting first or by using heapify:

def __init__(self, k: int, nums: List[int]):
    self.k = k
    # Inefficient: sorting entire array
    nums.sort(reverse=True)
    self.min_heap = nums[:k]
    heapq.heapify(self.min_heap)

Why This is Problematic:

• Sorting takes O(n log n) time, which is worse than O(n log k)
• Taking the first k elements after sorting in descending order gives you a max heap structure, not a min heap
• This breaks the invariant that the smallest of the k largest should be at the root

The Correct Approach: Use the add method for each element during initialization, which properly maintains the heap invariant and handles all edge cases consistently.