Leetcode 703. Kth Largest Element in a Stream

Problem Description

The problem requires designing a class, KthLargest, that handles a stream of integers and can return the kth largest element in the stream at any given point. The class constructor takes two arguments: an integer k denoting the rank of the largest element we want to find, and an integer array nums containing the initial elements of the stream.

The class has a method add that takes an integer as an argument, adds it to the stream, and then returns the kth largest element in the stream. It should be noted that it's about the kth largest element in sorted order, not the kth distinct element.

Example

Here is an example to illustrate how this class should function:

1    ```

java int k = 3; int[] arr = {4,5,8,2}; KthLargest kthLargest = new KthLargest(3, arr); kthLargest.add(3); // returns 4 kthLargest.add(5); // returns 5 kthLargest.add(10); // returns 5 kthLargest.add(9); // returns 8 kthLargest.add(4); // returns 8 ```

After initialization, the sorted stream is [2, 3, 4, 5, 8], and the 3rd largest number is 4. After adding number 3, the new stream is [2,3,4,5,8,3], and the 3rd largest number is still 4. And so on.

Approach

The ideal data structure for this problem is a min-heap. A min-heap always bubbles up the smallest element to the top, which in this case helps to maintain the kth largest elements in the heap at all times.

If the size of the heap is less than k, we just add the new number into it. If the size is equal to k, we check if the new number is larger than the top element (which is the smallest in the heap). If it is, we remove the top element and add the new number. Otherwise, we skip it because it’s less than kth largest element.

Implementation in different languages

C++

1
2cpp
3#include<queue>
4#include<vector>
5class KthLargest {
6public:
7    KthLargest(int k, vector<int>& nums) {
8        limit = k;
9        for(int i=0;i<nums.size();i++){
10            add(nums[i]);
11        }
12    }
13    
14    int add(int val) {
15        if(pq.size() < limit){
16            pq.push(val);
17        }
18        else if(val > pq.top()){
19            pq.pop();
20            pq.push(val);
21        }
22        return pq.top();
23    }
24
25private:
26    priority_queue<int,vector<int>,greater<int>> pq;
27    int limit;
28};

Java

1
2java
3class KthLargest {
4    final PriorityQueue<Integer> q;
5    final int k;
6
7    public KthLargest(int k, int[] a) {
8        this.k = k;
9        q = new PriorityQueue<>(k);
10        for (int n : a)
11            add(n);
12    }
13
14    public int add(int n) {
15        if (q.size() < k)
16            q.offer(n);
17        else if (q.peek() < n) {
18            q.poll();
19            q.offer(n);
20        }
21        return q.peek();
22    }
23}

Python

1
2python
3import heapq
4
5class KthLargest:
6    def __init__(self, k, nums):
7        self.heap = nums
8        self.k = k
9        heapq.heapify(self.heap)
10        while len(self.heap) > k:
11            heapq.heappop(self.heap)
12
13    def add(self, val):
14        if len(self.heap) < self.k:
15            heapq.heappush(self.heap, val)
16        elif val > self.heap[0]:
17            heapq.heapreplace(self.heap, val)
18        return self.heap[0]

JavaScript

1
2javascript
3class KthLargest {
4    constructor(k, nums) {
5        this.k = k;
6        this.minHeap = new MinHeap();
7        for(let num of nums){
8            this.add(num);
9        }
10    }
11    
12    add(val) {
13        if(this.minHeap.size() < this.k){
14            this.minHeap.offer(val);
15        }else if(val > this.minHeap.peek()){
16            this.minHeap.poll();
17            this.minHeap.offer(val);
18        }
19        return this.minHeap.peek();
20    }
21};
22
23// Construct a min heap
24class MinHeap {
25  constructor() {
26    this.heap = [];
27  }
28  
29  offer(element) {
30    this.heap.push(element);
31    this.bubbleUp();
32  }
33  
34  poll() {
35    let result = this.heap[0];
36    let end = this.heap.pop();
37    if (this.heap.length > 0) {
38      this.heap[0] = end;
39      this.bubbleDown();
40    }
41    return result;
42  }
43  
44  bubbleUp() { }
45  bubbleDown() { }
46  peek() { return this.heap[0]; }
47  size() { return this.heap.length; }
48}

C#

1
2csharp
3public class KthLargest {
4
5    private SortedList<int, int> nums;
6    private int maxSize;
7	
8    public KthLargest(int k, int[] arr) {
9        nums = new SortedList<int, int>(new DuplicateKeyComparer<int>());
10        maxSize = k;
11
12        for (int i = 0; i < arr.Length; i++) {
13            Add(arr[i]);
14        }
15    }
16	
17    public int Add(int val) {
18        nums.Add(val, 0);
19        if (nums.Count > maxSize) {
20            nums.RemoveAt(0);
21        }
22			
23        return nums.Keys[0];
24    }
25}
26
27public class DuplicateKeyComparer<TKey> : IComparer<TKey> where TKey : IComparable
28{
29    public int Compare(TKey x, TKey y)
30    {
31        int result = x.CompareTo(y);
32
33        if (result == 0)
34            return 1;   
35        else
36            return result;
37    }
38}

Ruby

1
2ruby
3class KthLargest
4    def initialize(k, nums)
5        @k = k
6        @nums = nums.sort.last(k)
7    end
8    
9    def add(val)
10        idx = (0...@nums.size).bsearch { |i| @nums[i] > val } || @nums.size
11        @nums.insert(idx, val)
12
13        @nums.pop if @nums.size > @k
14        @nums.last
15    end
16end

Go

1
2go
3import "container/heap"
4
5type minHeap []int
6
7func (h minHeap) Len() int {
8    return len(h)
9}
10
11func (h minHeap) Less(i, j int) bool {
12    return h[i] < h[j]
13}
14
15func (h minHeap) Swap(i, j int) {
16    h[i], h[j] = h[j], h[i]
17}
18
19func (h *minHeap) Push(x interface{}) {
20    *h = append(*h, x.(int))
21}
22
23func (h *minHeap) Pop() interface{} {
24    old := *h
25    n := len(old)
26    x := old[n-1]
27    *h = old[0 : n-1]
28    return x
29}
30
31type KthLargest struct {
32    nums minHeap
33    k    int
34}
35
36func Constructor(k int, nums []int) KthLargest {
37    h := &KthLargest{k: k, nums: nums}
38    heap.Init(&(h.nums))
39    for len(h.nums) > k {
40        heap.Pop(&(h.nums))
41    }
42    return *h
43}
44
45func (this *KthLargest) Add(val int) int {
46    heap.Push(&(this.nums), val)
47    for len(this.nums) > this.k {
48        heap.Pop(&(this.nums))
49    }
50    return this.nums[0]
51}

As you can see, the process of maintaining the kth largest number in a stream of numbers is pretty straightforward in any programming language. The main tool necessary to solve this problem efficiently is the min-heap data structure. Some languages provide built-in support for this type of data structure, while others require the explicit implementation of several methods to simulate a min-heap.