Solution: Kth Largest Element in a Stream

Let's solve the Kth Largest Element in a Stream problem using the Top K Elements pattern.

We'll cover the following

Statement
Solution
- Naive approach
- Optimized approach using Top K Elements

Statement

Given an infinite stream of integers (sorted or unsorted), nums, design a class to find the $k^{th}$ largest element in a stream.

Note: It is the $k^{th}$ largest element in the sorted order, not the $k^{th}$ distinct element.

The class should have the following functions, inputs, and return values:

Init(nums, k): It takes an array of integers nums and an integer k and initializes the class object.
Add(value): It takes one integer value, appends it to the stream, and returns the element representing the $k^{th}$ largest element in the stream.

Constraints:

$1 \leq k \leq 10^3$
$0 \leq$ nums.length $\leq 10^3$
$-10^3 \leq$ nums[i] $\leq 10^3$
$-10^3 \leq$ value $\leq 10^3$
At most, $10^3$ 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 $k^{th}$ element.

Solution

So far, you have probably brainstormed some approaches and have an idea of how to solve this problem. Let’s explore some of these approaches and figure out which one to follow based on considerations such as time complexity and any implementation constraints.

Naive approach

The naive solution is to first sort the data and then find the $k^{th}$ largest element. Insertion sort is an algorithm that can be used to sort the data as it appears. However, it also requires shifting the elements, greater than the inserted number, one place forward.

The overall time complexity of the algorithm becomes $O(n^2)$ , where $n$ is the number of elements in the data stream. The time complexity of each insertion is $O(n)$ and finding the $k^{th}$ largest element would take $O(1)$ time, assuming we are storing the data in an array. The space complexity is $O(1)$ .

Optimized approach using Top K Elements

To efficiently find the $k^{th}$ largest element in a stream of numbers, we use a min-heap that holds the top $k$ largest elements. This way, we don’t have to sort the entire list each time a new number is added. The $k^{th}$ largest element will change as new members come in, so we need a class to handle these dynamic updates.

With its ability to hold k elements, the min-heap ensures that the $k^{th}$ largest number is always at the root. We do this by adding new elements to the heap and removing the smallest one if the heap grows beyond k elements. This approach allows us quick access to the $k^{th}$ largest number, making the min-heap the most efficient tool for the job.

The slides below illustrate the core ideas of our algorithm:

Press + to interact

import java.util.*;
class KthLargest {
    public PriorityQueue<Integer> topKHeap;
    public int k;
    // Constructor to initialize heap and add values in it
    public KthLargest(int k, int[] nums) {
        System.out.println("\n\tInitializing the heap");
        this.k = k;
        topKHeap = new PriorityQueue<>();
        System.out.println("\tk = " + this.k);
        System.out.println("\tHeap: " + topKHeap);
        for (int element : nums) {
            System.out.println("Adding element " + element + " to the heap.");
            add(element);
        }
    }
    // Adds element in the heap and return the Kth largest
    public int add(int val) {
        // Add your implementation here
        return -1; // Placeholder return value
    }
    // Driver code
    public static void main(String[] args) {
        int[] nums = {3, 6, 9, 10};
        System.out.print("Initial stream: ");
        printArray(nums);
        KthLargest kLargest = new KthLargest(3, nums);
    }
    public static void printArray(int[] arr) {
        System.out.print("[");
        for (int i = 0; i < arr.length; i++) {
            System.out.print(arr[i]);
            if (i != arr.length - 1) {
                System.out.print(", ");
            }
        }
        System.out.print("]");
    }
}

Kth Largest Element in a Stream

Next, we’ll implement the add function, which adds an incoming number to the stream. If the heap size is less than $k$ , we’ll simply push the value into the heap. Otherwise, if the incoming value is greater than the smallest element (top of the heap), we’ll perform a pop operation to remove the smallest element and then push the new value into the heap. This ensures that the heap maintains only the k largest elements. Since we’re working with a min-heap, the $k^{th}$ largest element will always be at the top of the heap.

import java.util.*;
class KthLargest {
    public PriorityQueue<Integer> topKHeap;
    public int k;
    // Constructor to initialize heap and add values in it
    public KthLargest(int k, int[] nums) {
        System.out.println("\tInitializing the heap");
        this.k = k;
        topKHeap = new PriorityQueue<>();
        System.out.println("\tk = " + this.k);
        System.out.println("\tHeap: " + topKHeap);
        for (int element : nums) {
            System.out.println("\tAdding element " + element + " to the heap.");
            add(element);
        }
    }
    // Adds element in the heap and return the Kth largest
    public int add(int val) {
        if (topKHeap.size() < k) {
            System.out.println("\tSize of the heap is less than k, so we’ll push " + val);
            topKHeap.offer(val);
        } else if (val > topKHeap.peek()) {
            System.out.println("\tSize of the heap is equal to k and " + val + " is greater than the top of heap " + topKHeap.peek());
            System.out.println("\tWe will poll the top element from heap and offer the new element.");
            topKHeap.poll();
            topKHeap.offer(val);
        } else {
            System.out.println("\tSize of the heap is equal to k and " + val + " is equal to or smaller than the top of heap.");
            System.out.println("\tNo operation will be performed.");
        }
        System.out.println("\tHeap: " + topKHeap);
		return -1; // Placeholder return value
    }
    // Driver code
    public static void main(String[] args) {
        int[] nums = {3, 6, 9, 10};
        int[] temp = {3, 6, 9, 10};
        System.out.print("Initial stream: ");
        printArray(nums);
        System.out.println("\nk: " + 3);
        KthLargest kLargest = new KthLargest(3, nums);
        int[] val = {4, 7, 10, 8, 15};
        for (int i = 0; i < val.length; i++) {
            System.out.println("\tAdding a new number " + val[i] + " to the stream");
            temp = Arrays.copyOf(temp, temp.length + 1);
            temp[temp.length - 1] = val[i];
            System.out.print("\tNumber stream: ");
            printArray(temp);
			System.out.println("");
            kLargest.add(val[i]);
            System.out.println(new String(new char[100]).replace('\0', '-'));
        }
    }
    public static void printArray(int[] arr) {
        System.out.print("[");
        for (int i = 0; i < arr.length; i++) {
            System.out.print(arr[i]);
            if (i != arr.length - 1) {
                System.out.print(", ");
            }
        }
        System.out.print("]");
    }
}

Kth Largest Element in a Stream

import java.util.*;
class KthLargest {
    public PriorityQueue<Integer> topKHeap;
    public int k;
	// Constructor to initialize heap and add values in it
    public KthLargest(int k, int[] nums) {
        System.out.println("\tInitializing the heap");
        this.k = k;
        topKHeap = new PriorityQueue<>();
        System.out.println("\tk = " + this.k);
        System.out.println("\tHeap: " + topKHeap);
        for (int element : nums) {
            System.out.println("\tAdding element " + element + " to the heap.");
            add(element);
        }
    }
	// Adds element in the heap and return the Kth largest
    public int add(int val) {
        if (topKHeap.size() < k) {
            System.out.println("\n\tSize of the heap is less than k, so we’ll push " + val);
            topKHeap.offer(val);
        } else if (val > topKHeap.peek()) {
            System.out.println("\n\tSize of the heap is equal to k and " + val + " is greater than the top of heap " + topKHeap.peek());
            System.out.println("\tWe will poll the top element from heap and offer the new element.");
            topKHeap.poll();
            topKHeap.offer(val);
        } else {
            System.out.println("\n\tSize of the heap is equal to k and " + val + " is equal to or smaller than the top of heap.");
            System.out.println("\tNo operation will be performed.");
        }
        System.out.println("\tHeap: " + topKHeap);
        return topKHeap.peek();
    }
	// Driver code
    public static void main(String[] args) {
        int[] nums = {3, 6, 9, 10};
        int[] temp = {3, 6, 9, 10};
        System.out.print("Initial stream: ");
        printArray(nums);
        System.out.println("\nk: " + 3);
        KthLargest kLargest = new KthLargest(3, nums);
        int[] val = {4, 7, 10, 8, 15};
        for (int i = 0; i < val.length; i++) {
            System.out.println("\tAdding a new number " + val[i] + " to the stream");
            temp = Arrays.copyOf(temp, temp.length + 1);
            temp[temp.length - 1] = val[i];
            System.out.print("\tNumber stream: ");
            printArray(temp);
            System.out.println("\tKth largest element in the stream: " + kLargest.add(val[i]));
            System.out.println(new String(new char[100]).replace('\0', '-'));
        }
    }
    public static void printArray(int[] arr) {
        System.out.print("[");
        for (int i = 0; i < arr.length; i++) {
            System.out.print(arr[i]);
            if (i != arr.length - 1) {
                System.out.print(", ");
            }
        }
        System.out.print("]");
    }
}

Kth Largest Element in a Stream

import java.util.*;
class KthLargest {
    public PriorityQueue<Integer> topKHeap;
    public int k;
	// Constructor to initialize heap and add values in it
    public KthLargest(int k, int[] nums) {
        this.k = k;
        topKHeap = new PriorityQueue<>();
        for (int element : nums) {
            add(element);
        }
    }
	// Adds element in the heap and return the Kth largest
    public int add(int val) {
        if (topKHeap.size() < k) {
            topKHeap.offer(val);
        } else if (val > topKHeap.peek()) {
            topKHeap.poll();
            topKHeap.offer(val);
        }
        return topKHeap.peek();
    }
    public static void main(String[] args) {
        int[] nums = {3, 6, 9, 10};
        int[] temp = {3, 6, 9, 10};
        System.out.print("Initial stream: ");
        printArray(nums);
        System.out.println("\nk: " + 3);
        KthLargest kLargest = new KthLargest(3, nums);
        int[] val = {4, 7, 10, 8, 15};
        for (int i = 0; i < val.length; i++) {
            System.out.println("\tAdding a new number " + val[i] + " to the stream");
            temp = Arrays.copyOf(temp, temp.length + 1);
            temp[temp.length - 1] = val[i];
            System.out.print("\tNumber stream: ");
            printArray(temp);
            System.out.println("\n\tKth largest element in the stream: " + kLargest.add(val[i]));
            System.out.println(new String(new char[100]).replace('\0', '-'));
        }
    }
	public static void printArray(int[] arr) {
		System.out.print("[");
        for (int i = 0; i < arr.length; i++) {
            System.out.print(arr[i]);
            if (i != arr.length - 1) {
                System.out.print(", ");
            }
        }
		System.out.print("]");
    }
}

Kth Largest Element in a Stream

Solution summary

To recap, the solution to this problem can be divided into the following parts:

Initialize a min-heap in the constructor to store the elements. Iterate through the elements in nums and call the add function for each element.
In the add function, if the size of the heap is less than $k$ , push the number into the heap. Otherwise, if the incoming value is greater than the smallest element (top of the heap), perform a pop operation to remove the smallest element, and then push the new value into the heap.
After adding all the numbers, the $k^{th}$ largest element can be found at the root of the heap.

Time complexity

Init(nums, k): In the constructor, we iterate through the nums list with $n$ elements and add them to the heap. Adding each element to the heap takes, at most, $O(\log k)$ time. Therefore, the overall time complexity of the constructor is $O(n \log k)$ .
Add(value): In the add function, we perform operations like pushing and popping elements from the heap, which take, at most, $O(\log k)$ time. Therefore, the time complexity of the add function is $O(\log k)$ .

Space complexity

The space complexity will be $O(k)$ where $k$ is the $k^{th}$ largest element we need to find in a stream. This is because the heap is initialized to have a maximum size of $k$ , resulting in an overall space complexity of $O(k)$ .

Level up your interview prep. Join Educative to access 70+ hands-on prep courses.

Getting Started

Two Pointers

Fast and Slow Pointers

Sliding Window

Merge Intervals

In-Place Manipulation of a Linked List

Two Heaps

K-way merge

Top K Elements

Modified Binary Search

Subsets

Greedy Techniques

Backtracking

Dynamic Programming

Cyclic Sort

Topological Sort

Sort and Search

Matrices

Stacks

Graphs

Tree Depth-First Search

Tree Breadth-First Search

Trie

Hash Maps

Knowing What to Track

Union Find

Custom Data Structures

Bitwise Manipulation

Challenge Yourself

Conclusion

Solution: Kth Largest Element in a Stream

Statement

Solution

Naive approach

Optimized approach using Top K Elements

Step-by-step solution construction

Just the code

Solution summary

Time complexity

Space complexity