Skip to main content

Heaps

Kth Largest Element in a Stream Solution In C++/Python/Java/JS

Problem Description

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 kth largest element in the pool of test scores so far.

Example

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


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 <= 104
  • 1 <= k <= nums.length + 1
  • -104 <= nums[i] <= 104
  • -104 <= val <= 104
  • At most 104 calls will be made to add.
💡
Try it Yourself First
Imagine a stock market analyst who wants to track the k-th highest stock price in real time as new prices are recorded. This helps them dynamically identify significant price trends. To achieve this, we'll create a class KthLargest that returns the k-th highest stock price from an incoming stream of prices.
Specifically, we need to implement:
The constructor KthLargest(int k, int[] prices), which initializes the class with k and the initial list of stock prices.
The function add(int price), which adds a new stock price to the existing stream and returns the k-th highest price in the updated data.

Brute Force Approach

Intuition

Our goal is to efficiently maintain and retrieve the k-th largest element from a growing stream of numbers. Instead of storing numbers unsorted and sorting after each addition, we ensure that the list remains sorted at all times, allowing us to fetch the k-th largest element instantly.
To achieve this, we store all numbers in a sorted list and use binary search for efficient insertion. This prevents unnecessary full re-sorting and speeds up updates.
We begin by initializing the list with the given array nums and sorting it in ascending order. This ensures that the k-th largest element is always at index stream.length - k.
When a new number val is added through add(int val), we determine its correct position in the sorted list using binary search:

  • Start with the entire list as the search space.
  • Compare val with the middle element stream[mid].
  • If stream[mid] == val, insert val at index mid.
  • If stream[mid] < val, continue searching in the right half.
  • If stream[mid] > val, search in the left half.
  • Once the correct position is found, insert val while maintaining the sorted order.

Since retrieving the k-th largest element is simply an index lookup at stream[stream.length - k], this approach significantly improves efficiency. Using binary search for insertion reduces the time complexity to O(log n) for each addition, making this method much faster than naive sorting-based solutions.

Approach

Step 1: Initialize Class Variables

  • Store k as a class variable to track the required k-th largest element.
  • Maintain a sorted list stream to store all numbers in ascending order.
  • Add all elements from nums to stream and sort it.

Step 2: Implement the add(int val) Function

  • Call the helper function getIndex(val) to find the correct position to insert val.
  • Insert val at the determined index in stream, ensuring it remains sorted.
  • Return the k-th largest element, found at stream[stream.length - k].

Step 3: Implement the getIndex(int val) Function

  • Define the search space with left = 0 and right = stream.length - 1.
  • While left <= right:
  • Compute the middle index mid = (left + right) / 2.
  • Compare val with stream[mid]:
  • If val == stream[mid], return mid.
  • If val < stream[mid], move search to the left half (right = mid - 1).
  • If val > stream[mid], move search to the right half (left = mid + 1).
  • Once the correct index is found, return it for insertion.

Dry Run

Input:

  • Operations:
  • ["KthLargest", "add", "add", "add", "add", "add"]
  • Arguments:
  • [[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
  • Expected Output:
  • [null, 4, 5, 5, 8, 8]

Step 1: Initialize KthLargest

  • Input: KthLargest(3, [4, 5, 8, 2])
  • Sorted stream: [2, 4, 5, 8]
  • Since k = 3, the k-th largest element is at index stream.length - k = 4 - 3 = 1.
  • Current k-th largest element: 4
  • Output: null (constructor does not return a value)

Step 2: Process add(3)

  • Insert 3 into sorted stream using binary search:
  • Updated stream: [2, 3, 4, 5, 8]
  • k-th largest element (index 5 - 3 = 2): 4
  • Output: 4

Step 3: Process add(5)

  • Insert 5 into sorted stream using binary search:
  • Updated stream: [2, 3, 4, 5, 5, 8]
  • k-th largest element (index 6 - 3 = 3): 5
  • Output: 5

Step 4: Process add(10)

  • Insert 10 into sorted stream using binary search:
  • Updated stream: [2, 3, 4, 5, 5, 8, 10]
  • k-th largest element (index 7 - 3 = 4): 5
  • Output: 5

Step 5: Process add(9)

  • Insert 9 into sorted stream using binary search:
  • Updated stream: [2, 3, 4, 5, 5, 8, 9, 10]
  • k-th largest element (index 8 - 3 = 5): 8
  • Output: 8

Step 6: Process add(4)

  • Insert 4 into sorted stream using binary search:
  • Updated stream: [2, 3, 4, 4, 5, 5, 8, 9, 10]
  • k-th largest element (index 9 - 3 = 6): 8
  • Output: 8

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

Code for All Languages
C++ 
// Approach: Binary Search Insertion with Sorted List  
// Language: C++  

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest* obj = new KthLargest(k, nums);
 * int param_1 = obj->add(val);
 */

class KthLargest {
public:
    
    // Stores the sorted stream of numbers
    vector<int> stream;
    
    // Stores the k-th largest position to track
    int k;

    // Step 1: Initialize the object with k and the initial set of numbers
    KthLargest(int k, vector<int>& nums) {
        this->k = k;
        
        // Insert all numbers into the stream
        for (int num : nums) {
            stream.push_back(num);
        }
        
        // Sort the stream to maintain order
        sort(stream.begin(), stream.end());
    }

    // Step 2: Add a new number and return the k-th largest element
    int add(int val) {
        
        // Find the correct position using binary search
        int index = getIndex(val);
        
        // Insert value at the correct position while maintaining order
        stream.insert(stream.begin() + index, val);
        
        // Return the k-th largest element
        return stream[stream.size() - k];
    }

private:
    
    // Step 3: Find the correct insertion index using binary search
    int getIndex(int val) {
        
        // Define the search space
        int left = 0, right = stream.size() - 1;
        
        // Perform binary search to find the correct position
        while (left <= right) {
            
            // Calculate middle index
            int mid = (left + right) / 2;
            
            // Get the middle element
            int midValue = stream[mid];
            
            // If the middle element is equal to val, return its index
            if (midValue == val) return mid;
            
            // If val is smaller, search in the left half
            if (midValue > val) {
                right = mid - 1;
            } 
            
            // If val is greater, search in the right half
            else {
                left = mid + 1;
            }
        }
        
        // Return the correct position for insertion
        return left;
    }
};

Java 
// Approach: Binary Search Insertion with Sorted List  
// Language: Java  

import java.util.*;

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest obj = new KthLargest(k, nums);
 * int param_1 = obj.add(val);
 */
class KthLargest {
    
    // Stores the sorted list of numbers
    private List<Integer> stream;
    
    // Tracks the k-th largest element
    private int k;

    // Step 1: Initialize the object with k and the initial set of numbers
    public KthLargest(int k, int[] nums) {
        this.k = k;
        
        // Initialize the stream list with the given numbers
        stream = new ArrayList<>(nums.length);
        
        // Add elements from nums to the stream list
        for (int num : nums) {
            stream.add(num);
        }
        
        // Sort the stream to maintain order
        Collections.sort(stream);
    }

    // Step 2: Add a new number and return the k-th largest element
    public int add(int val) {
        
        // Find the correct position to insert the new value
        int index = getIndex(val);
        
        // Insert value at the correct position while maintaining sorted order
        stream.add(index, val);
        
        // Return k-th largest element from the stream
        return stream.get(stream.size() - k);
    }

    // Step 3: Find the correct insertion index using binary search
    private int getIndex(int val) {
        
        // Define the search space
        int left = 0, right = stream.size() - 1;
        
        // Perform binary search to find the correct position
        while (left <= right) {
            
            // Calculate middle index
            int mid = (left + right) / 2;
            
            // Get the middle element
            int midElement = stream.get(mid);
            
            // If the middle element is equal to val, return its index
            if (midElement == val) return mid;
            
            // If val is smaller, search in the left half
            if (midElement > val) {
                right = mid - 1;
            } 
            
            // If val is greater, search in the right half
            else {
                left = mid + 1;
            }
        }
        
        // Return the correct position for insertion
        return left;
    }
}

Python 
# Approach: Binary Search Insertion with Sorted List  
# Language: Python  

import bisect 
 
# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)

class KthLargest(object):

    # Step 1: Initialize the object with k and the initial set of numbers
    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        
        # Store k to track the k-th largest element
        self.k = k
        
        # Store the sorted list of numbers
        self.stream = sorted(nums)

    # Step 2: Add a new number and return the k-th largest element
    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        
        # Find the correct position using binary search
        index = bisect.bisect_left(self.stream, val)
        
        # Insert the value in the correct position to maintain order
        self.stream.insert(index, val)
        
        # Return the k-th largest element
        return self.stream[-self.k]

Javascript 
// Approach: Binary Search Insertion with Sorted List  
// Language: JavaScript  

/**
 * @param {number} k
 * @param {number[]} nums
 */
var KthLargest = function(k, nums) {
    
    // Store k to track the k-th largest element
    this.k = k;
    
    // Store the sorted list of numbers
    this.stream = nums.sort((a, b) => a - b);
};

/** 
 * @param {number} val
 * @return {number}
 */
KthLargest.prototype.add = function(val) {
    
    // Step 1: Find the correct position using binary search
    let left = 0, right = this.stream.length - 1;
    
    while (left <= right) {
        
        // Calculate middle index
        let mid = Math.floor((left + right) / 2);
        
        // If val is greater, search in the right half
        if (this.stream[mid] < val) {
            left = mid + 1;
        } 
        
        // If val is smaller, search in the left half
        else {
            right = mid - 1;
        }
    }
    
    // Step 2: Insert value in the correct position to maintain sorted order
    this.stream.splice(left, 0, val);
    
    // Step 3: Return the k-th largest element
    return this.stream[this.stream.length - this.k];
};

/** 
 * Your KthLargest object will be instantiated and called as such:
 * var obj = new KthLargest(k, nums)
 * var param_1 = obj.add(val)
 */

Time Complexity Analysis:O(N^2+N⋅M)

Let M be the size of the initial stream nums given in the constructor. Let N be the number of calls of add.

Time Complexity: O(N^2+N⋅M)

  • The constructor involves creating a list stream from nums, which takes O(M) time. Then, sorting this list takes O(M⋅logM) time. Thus, the time complexity of the constructor is O(M⋅logM) time.
  • The add function involves running a binary search on stream. Because the total size of stream at the end would be O(M+N), each binary search is bounded by a time complexity of O(log(M+N)). Moreover, adding a number in stream can take worst-case O(M+N) time, as adding an element in the middle of a list can offset all the elements to its right. Then, the time complexity of a single add call would be O(M+N+log(M+N)). Because add is called N times, the time complexity of all the add calls would be O(N⋅(M+N+log(M+N))).
  • We see that after expanding the time complexity for the add function, the N⋅M and N2 terms dominate all the other log terms in our calculations, so the total time complexity is O(N2+N⋅M)

Final Time Complexity: O(N^2+N⋅M)

Space Complexity Analysis: O(M + N)

Sorting and Storage (O(M + N))

  • The stream list stores all elements from the initial nums array (N) and the incoming add() calls (M), requiring O(M + N) space.

Auxiliary Space (O(1))

  • Binary search runs in-place and does not use extra memory.
  • Insertion modifies the existing list rather than creating a new one.
  • No additional data structures (like heaps or extra arrays) are used, keeping auxiliary space at O(1).

Final Space Complexity:O(M + N)

  • The total space required is due to maintaining the sorted list, resulting in O(M + N) overall space complexity.
The initial approach maintained a fully sorted list, leading to O(n log n) complexity for updates. To optimize, we use a min-heap that efficiently tracks only the top k elements, reducing insertion time to O(log k). Instead of sorting the entire stream, we dynamically update the heap, ensuring the k-th largest element is always at the root, making retrieval O(1). This transition significantly improves both time and space efficiency.

Optimal Approach

Intuition

The initial approach of maintaining a fully sorted list for tracking the k-th largest element becomes inefficient as the stream grows. Sorting after every insertion leads to O(n log n) complexity, which is unnecessary since we only care about the top k elements rather than the entire stream. To optimize, we use a min-heap (priority queue), which dynamically maintains only the k largest elements. The key insight is that the smallest element among these k elements represents the k-th largest element in the entire stream. Imagine a scenario where k = 3 and the current stream is [0, 4, 6, 9].

  • Before adding a new value, the k-th largest element is 4.
  • If the incoming value is 2, it is smaller than 4, so the top 3 largest numbers remain unchanged, and we can immediately return 4.
  • If the incoming value is 7, which is greater than both 4 and 6, it replaces 4 in the top 3 elements, making 6 the new k-th largest element, while 4 is removed.

This example highlights that only the top k numbers affect the result, and tracking unnecessary elements increases computational overhead. A min-heap efficiently solves this by keeping only the top k elements and ensuring that the smallest of them (k-th largest overall) is at the top of the heap. If a new number is greater than this smallest value, we insert it and remove the smallest element, keeping the heap size constant. By leveraging a min-heap, we avoid sorting large lists and reduce the update time to O(log k) per insertion, significantly improving efficiency while maintaining quick access to the k-th largest element.

Approach

Step 1: Initialize Class Variables
*Store k as a class variable to track the required k-th largest element.

  • Use a min heap (priority queue) to efficiently maintain only the top k largest elements.

Step 2: Process Initial Stream

  • Iterate through each number in the given nums array.
  • Call the add(num) function to insert each element into the heap while ensuring it contains only the k largest elements.

Step 3: Insert New Element in add(val)

  • If val is greater than the smallest element in the heap or the heap has fewer than k elements, insert val into the heap.
  • If the heap size exceeds k, remove the smallest element to keep only the top k values.

Step 4: Return the k-th Largest Element

  • The smallest element in the heap (root) is the k-th largest element.
  • Return the top of the heap as the current k-th largest element.

Dry Run

Input:

  • Operations:
  • ["KthLargest", "add", "add", "add", "add", "add"]
  • Arguments: [[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
  • Expected Output: [null, 4, 5, 5, 8, 8]

Step 1: Initialize Min Heap

  • Input: KthLargest(3, [4, 5, 8, 2])
  • Min Heap: [] (to store the k largest elements)
  • Insert elements one by one and maintain only k elements:
  • Add 4 → Heap: [4]
  • Add 5 → Heap: [4, 5]
  • Add 8 → Heap: [4, 5, 8]
  • Add 2 → Heap: [2, 4, 8, 5] (Heap size exceeds k, remove 2)
  • Final Heap: [4, 5, 8]
  • k-th Largest Element: 4 (smallest in the heap)
  • Output: null (constructor does not return a value)

Step 2: Process add(3)

  • New Value: 3
  • Since 3 < 4 (k-th largest), heap remains unchanged.
  • Return k-th largest: 4
  • Heap: [4, 5, 8]
  • Output: 4

Step 3: Process add(5)

  • New Value: 5
  • Since 5 > 4 (k-th largest), add 5 to heap.
  • Heap: [4, 5, 8, 5]
  • Heap size exceeds k, remove 4 (smallest).
  • Final Heap: [5, 5, 8]
  • Return k-th largest: 5
  • Output: 5

Step 4: Process add(10)

  • New Value: 10
  • Since 10 > 5 (k-th largest), add 10 to heap.
  • Heap: [5, 5, 8, 10]
  • Heap size exceeds k, remove 5 (smallest).
  • Final Heap: [5, 10, 8]
  • Return k-th largest: 5
  • Output: 5

Step 5: Process add(9)

  • New Value: 9
  • Since 9 > 5 (k-th largest), add 9 to heap.
  • Heap: [5, 9, 8, 10]
  • Heap size exceeds k, remove 5 (smallest).
  • Final Heap: [8, 9, 10]
  • Return k-th largest: 8
  • Output: 8

Step 6: Process add(4)

  • New Value: 4
  • Since 4 < 8 (k-th largest), heap remains unchanged.
  • Return k-th largest: 8
  • Heap: [8, 9, 10]
  • Output: 8

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

Code for All Languages
C++ 
// Approach: Min Heap (Priority Queue) to Maintain k Largest Elements  
// Language: C++  

#include <iostream>
#include <vector>
#include <queue>
using namespace std;

class KthLargest {
private:
    
    // Min Heap to store the top k largest elements
    priority_queue<int, vector<int>, greater<int>> minHeap;
    
    // Stores the value of k
    int k;

public:
    
    // Step 1: Initialize k and process the initial stream
    KthLargest(int k, vector<int>& nums) {
        this->k = k;

        // Insert elements into the heap while maintaining size k
        for (int num : nums) {
            add(num);
        }
    }

    // Step 2: Add a new number and return the k-th largest element
    int add(int val) {
        
        // Add to minHeap if we haven't processed k elements yet
        // or if val is greater than the k-th largest element (heap top)
        if (minHeap.size() < k || minHeap.top() < val) {
            minHeap.push(val);
            
            // If heap size exceeds k, remove the smallest element
            if (minHeap.size() > k) {
                minHeap.pop();
            }
        }
        
        // Return the k-th largest element (top of the minHeap)
        return minHeap.top();
    }
};

Java 
// Approach: Min Heap (Priority Queue) to Maintain k Largest Elements  
// Language: Java  

import java.util.*;

class KthLargest {
    
    // Min Heap to store the top k largest elements
    private PriorityQueue<Integer> minHeap;
    
    // Stores the value of k
    private int k;

    // Step 1: Initialize k and process the initial stream
    public KthLargest(int k, int[] nums) {
        this.k = k;
        
        // Initialize the minHeap
        minHeap = new PriorityQueue<>();
        
        // Insert elements into the heap while maintaining size k
        for (int num : nums) {
            add(num);
        }
    }

    // Step 2: Add a new number and return the k-th largest element
    public int add(int val) {
        
        // Add to minHeap if we haven't processed k elements yet
        // or if val is greater than the k-th largest element (heap top)
        if (minHeap.size() < k || minHeap.peek() < val) {
            minHeap.offer(val);
            
            // If heap size exceeds k, remove the smallest element
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }
        
        // Return the k-th largest element (top of the minHeap)
        return minHeap.peek();
    }
}

Python 
# Approach: Min Heap (Priority Queue) to Maintain k Largest Elements  
# Language: Python  

import heapq  

class KthLargest:
    
    # Step 1: Initialize k and process the initial stream
    def __init__(self, k, nums):
        """
        :type k: int
        :type nums: List[int]
        """
        
        # Store k to track the k-th largest element
        self.k = k
        
        # Initialize a min heap
        self.minHeap = []
        
        # Insert elements into the heap while maintaining size k
        for num in nums:
            self.add(num)

    # Step 2: Add a new number and return the k-th largest element
    def add(self, val):
        """
        :type val: int
        :rtype: int
        """
        
        # Add to minHeap if we haven't processed k elements yet
        # or if val is greater than the k-th largest element (heap top)
        if len(self.minHeap) < self.k or self.minHeap[0] < val:
            heapq.heappush(self.minHeap, val)
            
            # If heap size exceeds k, remove the smallest element
            if len(self.minHeap) > self.k:
                heapq.heappop(self.minHeap)
        
        # Return the k-th largest element (top of the minHeap)
        return self.minHeap[0]

Javascript 
// Approach: Min Heap (Priority Queue) to Maintain k Largest Elements  
// Language: JavaScript  

class KthLargest {
    
    // Step 1: Initialize k and process the initial stream
    constructor(k, nums) {
        
        // Store k to track the k-th largest element
        this.k = k;
        
        // Initialize a min heap (priority queue)
        this.minHeap = [];
        
        // Insert elements into the heap while maintaining size k
        nums.forEach(num => this.add(num));
    }

    // Step 2: Add a new number and return the k-th largest element
    add(val) {
        
        // Add to minHeap if we haven't processed k elements yet
        // or if val is greater than the k-th largest element (heap top)
        if (this.minHeap.length < this.k || this.minHeap[0] < val) {
            this._heapPush(val);
            
            // If heap size exceeds k, remove the smallest element
            if (this.minHeap.length > this.k) {
                this._heapPop();
            }
        }
        
        // Return the k-th largest element (top of the minHeap)
        return this.minHeap[0];
    }

    // Helper function to push an element into the min heap (maintains order)
    _heapPush(val) {
        this.minHeap.push(val);
        this.minHeap.sort((a, b) => a - b); // Ensure min-heap order
    }

    // Helper function to remove the smallest element from the heap
    _heapPop() {
        this.minHeap.shift(); // Remove the smallest element
    }
}

Time Complexity Analysis:O((M+N)⋅logk)

Let M be the size of the initial stream nums given in the constructor, and let N be the number of calls to add.

Time Complexity:O((M+N)⋅logk)

  • The add function involves adding and removing an element from a heap of size k, which is an O(logk) operation. Since the add function is called N times, the total time complexity for all add calls is O(N⋅logk).
  • The constructor also calls add M times to initialize the heap, leading to a time complexity of O(M⋅logk).

Final Time Complexity: O((M+N)⋅logk).

Space Complexity Analysis: O(k)

Min Heap Storage (O(k))

  • The min heap maintains at most k elements at any time.
  • Since we only track the k largest elements, the heap never grows beyond O(k).

Input Storage (O(n)) (not counted in auxiliary space complexity)

  • The input array nums is provided externally and does not count toward additional space usage.
  • We process nums one element at a time using the add() function, so we do not store the entire stream separately.

Auxiliary Space (O(1))

  • The heap operations (insertion and deletion) are performed in-place, requiring only a few extra variables (k, val).
  • No additional data structures or recursive calls are used.

Final Space Complexity:O(K)

  • The dominant factor is the min heap size, which is limited to k, making the overall space complexity O(k).

Similar Problems

Now we have successfully tackled this problem, let’s try these similar problems.

Kth Largest Element in an Array

Given an integer array nums and an integer k, return the kth largest element in the array.Note that it is the kth largest element in the sorted order, not the kth distinct element.

Kth Smallest Element in a Sorted Matrix

Given an n x n matrix where each of the rows and columns is sorted in ascending order, return the kth smallest element in the matrix.Note that it is the kth smallest element in the sorted order, not the kth distinct element.