Interview Problem Explained: Kth Largest Element in a Stream (Heap Technique)

Understanding the Problem: Kth Largest Element in a Stream

Every year, the admissions office at a university processes thousands of test scores submitted by applicants. Let’s imagine a system where, each time a new score comes in, it must instantly calculate the Kth highest score to dynamically update admission thresholds.

For example, let’s say K = 3, and the current scores are [90, 85]:

• When a new score 88 is added → [90, 88, 85] → the 3rd highest score is 85

• When 92 is added → [92, 90, 88, 85] → the 3rd highest score is 88

This problem challenges your ability to efficiently manage a stream of data and dynamically track the Kth largest value as new elements are added in real-time.

Problem Requirements:

1. Initialization: Build the system using an initial list of scores (nums) and an integer k.

2. Dynamic Updates: Each time the method add(val) is called, a new score is added to the stream, and the current Kth largest value must be returned.

3. Efficiency: The algorithm must scale to large volumes of data (e.g. tens of thousands of applicants) and respond instantly.

This is Leetcode Problem #703.
Leetcode’s problem descriptions are often brief and abstract.
Codetree explains them more deeply with intuitive context and visual walkthroughs to help you truly understand the logic behind the problem.

Step 1: Naive Approach

Idea

Your first instinct should be to implement exactly what the problem is asking for.

The requirement here is:
“We want to know the Kth largest value each time a new value is added.”
→ So, every time a new value is added to the stream, we sort the entire list in descending order and return the Kth largest element.

Logic Breakdown

1. Append the new value to the data array

2. Sort the array in descending order

3. Return data[k - 1]

Code Implementation

class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.data = nums[:]                 # Store all values as-is

    def add(self, val: int) -> int:
        self.data.append(val)               # ① Add new value
        self.data.sort(reverse=True)        # ② Sort in descending order
        return self.data[self.k - 1]        # ③ Return the Kth largest

Time and Space Complexity

Let N be the number of elements in the array and Q be the number of add calls.

• Each add operation takes O(N log N) due to sorting.

• Therefore, the total time complexity becomes O(Q × N log N).

This approach results in time limit issues when the dataset becomes large (e.g. N, Q ≤ 10⁴), and is not scalable for real-time stream processing.

Step 2. Sorted Insertion Optimization

Idea

As we saw in Step 1, sorting the entire array every time a new element is added leads to a time complexity of O(N log N), which is inefficient for large inputs. But if the array is already sorted, we can optimize by inserting the new element at the correct position while keeping the array sorted.

This way, we avoid sorting the full array again from scratch each time.

Logic Breakdown

1. Initially sort the nums array in descending order (only once)

2. Use linear search to find the correct index to insert the new value val

3. Insert the value at the correct position, shifting other elements to the right

4. Return the element at index k - 1 (the Kth largest)

This can be done using either an array or a linked list as the underlying data structure.

• Array: inserting at a specific index requires shifting elements → O(N)

• Linked list: inserting at a node is O(1), but you still need O(N) to find the position

Code Implementation

2‑1. Using an Array (List) + Linear Search

In Python lists, inserting an element requires shifting all the elements after the insertion point one position to the right.

class KthLargestArr:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        # Sort once in descending order
        self.data = sorted(nums, reverse=True)

    def add(self, val: int) -> int:
        # ① Find insertion index (O(N))
        idx = 0
        while idx < len(self.data) and self.data[idx] > val:
            idx += 1
        # ② Shift elements and insert
        self.data.append(0)  # Add dummy space
        for j in range(len(self.data) - 1, idx, -1):
            self.data[j] = self.data[j - 1]
        self.data[idx] = val
        return self.data[self.k - 1]

Time and Space Complexity

• Finding the insertion index: O(N)

• Shifting and inserting: O(N)

• Overall, each add call = O(N)

• Total complexity = O(Q × N)

2‑2. Using a Singly Linked List

Node-level insertion takes O(1) time, but since you have to traverse from the beginning to find the insertion position, the overall time complexity is O(N).

class ListNode:
    def __init__(self, val, nxt=None):
        self.val = val
        self.next = nxt

class KthLargestLL:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.head = ListNode(float('inf'))  # Dummy head with max value
        for n in sorted(nums, reverse=True):
            self._insert(n)

    def _insert(self, val):
        # Maintain descending order
        prev, cur = self.head, self.head.next
        while cur and cur.val > val:
            prev, cur = cur, cur.next
        prev.next = ListNode(val, cur)

    def add(self, val: int) -> int:
        self._insert(val)  # Find position and insert
        # Traverse to the Kth node
        cur = self.head.next
        for _ in range(self.k - 1):
            cur = cur.next
        return cur.val

Time and Space Complexity

• Insertion: O(N) (due to position search)

• Finding the Kth element: O(k) → included in O(N)

• Total time complexity per add: O(N)

Step 3. Using a Heap (Priority Queue)

Idea

In Step 2, we focused on optimizing time complexity by maintaining a sorted structure. However, even with sorted insertion, the time complexity remains O(N), which is still inefficient for large-scale streaming data.

Let’s take a step back and reconsider: What do we really need?

We only want the Kth largest value at any point. That means we don’t actually need to keep all the data sorted—we only care about the top K largest elements. So instead of managing unnecessary data or maintaining full sort order, we can optimize further by using a heap.
A min-heap of size k is a perfect fit because:

• It always keeps the smallest of the top K elements at the root

• That root is, by definition, the Kth largest overall

Logic Breakdown

Let’s break down how to use a heap for this problem.

1. Initialize the heap

   • Use the initial nums list to build a min-heap

   • If the heap size exceeds k, remove the smallest element (heap root)

2. Add a new value

   • If the heap size is less than k, push the new value directly

   • Otherwise, compare with the current root:

     • If the new value is larger → replace the root with the new value

3. Return the Kth largest

   • The root of the heap is always the Kth largest value

Code Implementation

import heapq

class KthLargest:
    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.heap = nums[:]
        heapq.heapify(self.heap)  # Convert list to a min-heap

        # Remove smallest elements if heap is too big
        while len(self.heap) > k:
            heapq.heappop(self.heap)

    def add(self, val: int) -> int:
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, val)
        elif val > self.heap[0]:  # Replace only if new val is bigger
            heapq.heapreplace(self.heap, val)
        return self.heap[0]

Time and Space Complexity

• Each add() call takes O(log k) time due to heap operations

• Overall time complexity for Q add() calls is O(Q × log k)

• Space complexity is O(k) since the heap stores at most k elements

This is the most efficient solution for this problem, especially when Q is large and real-time performance is critical.

Frequently Asked Questions

Q. Why do we use a min-heap?
Because a min-heap always keeps the smallest value at the top.
So if we maintain a min-heap of size k, the root node always represents the Kth largest element in the entire stream.

Q. What does "Kth largest" or "Kth smallest" element mean exactly?
It refers to the rank when all the elements are sorted.

For example, given [5, 1, 3, 2, 4], if we sort in descending order → [5, 4, 3, 2, 1]

• 1st largest = 5

• 2nd largest = 4

• 3rd largest = 3

If we sort in ascending order → [1, 2, 3, 4, 5]

• 1st smallest = 1

• 2nd smallest = 2

• 3rd smallest = 3

Q. How do we handle new elements in the stream?
When a new element is added, we need to insert it into our existing data structure in a way that still allows us to retrieve the Kth largest element efficiently. In earlier approaches, we added the value and then sorted the list, which is inefficient. We then improved this by inserting while maintaining sorted order.

Finally, we optimized even further by realizing that we only need to track the largest K values, and that’s how we arrived at using a min-heap.

Q. Can this be applied to other Top-K problems?
Be careful with this line of thinking. Just because a problem says “Top-K” doesn't mean a heap is always the best or only solution.

Instead, you should:

• First build a naive solution

• Then identify inefficiencies

• And only then consider if heap is a natural fit for optimizing the logic

It’s this iterative thought process—not memorizing patterns—that builds your real problem-solving skills.

How to Prepare for Interviews with the "Kth Largest element" Problem?

The “Kth Largest Element in a Stream” problem is a great practice example for dynamic data management and heap usage.
While it might not appear exactly as-is in real interviews, the process of exploring naive → optimized solutions is extremely valuable.

This problem helps you strengthen:

• Your understanding of the heap data structure and Python’s heapq module

• Your ability to maintain and update Top-K elements in real-time

• Your skill in analyzing and improving time and space complexity for streaming data

Similar Problems to Practice:

• Kth Largest Element in an Array

• Top K Frequent Elements

• Find Median from Data Stream (Codetree)

These problems all involve fixed-size Top-K logic or stream-based statistics, and they’re great for practicing heap-based thinking.

How to Review After Solving This Problem

1. Try implementing a heap from scratch

• Recreate a min-heap using a list without the heapq module, and understand how heappush and heappop work.

• While this isn’t typically required for coding tests, you should be able to clearly explain your implementation if asked during an interview.

2. Run time complexity experiments

• Test the performance difference between the sorting-based and heap-based approaches with real input data.

3. Take on modified versions of the problem

• Try solving “Kth Smallest in a Stream” by switching to a max-heap instead.

• In each problem involving heaps, start from the naive approach, then identify the specific inefficiency that the heap is designed to address.