Kth Largest Element From Data Stream Editorial

Practice Problem Link: Kth Largest Element From Data Stream

Please make sure to try solving the problem yourself before looking at the editorial.

Problem Statement

Given an incoming stream of data, you need to find the k^th largest element at each step.

Implement the KthLargest class:

• KthLargest(int k) initializes the KthLargest object with the integer k and is called at the beginning.
• int add(int num) adds the integer num from the data stream, and returns the k^th largest element until now.

int add(int num) returns -1 if there aren't k elements so far.

Naive Approach

We can solve this problem naively by sorting the current list, and printing the Kth element in the list, in each function call.

Analysis

• Time Complexity: O(n * logn)
• Space Complexity: O(n)

Implementation

C++

class KthLargest {
public:
    /** initialize your data structure here. */
	vector<int> elements;
	int K = 0;
    KthLargest(int k) {
		elements.clear();
		K = k;
    }
    
    int add(int num) { 
		elements.push_back(num);
		if(elements.size() < K) {
			return -1;
		}
		sort(elements.rbegin(), elements.rend());
		return elements[K - 1];
    }
};

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest* obj = new KthLargest(k);
 * int value = obj->addNum(num);
 */

Java

class KthLargest {
    /** initialize your data structure here. */
	ArrayList<Integer> elements;
	int k;
    public KthLargest(int k) {
		this.k = k;
        elements = new ArrayList<>();
    }
    
    public int add(int num) {
		elements.add(num);
		if(elements.size() < k) {
			return -1;
		}
        Collections.sort(elements);
		return elements.get(elements.size() - k);
    }
}

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest obj = new KthLargest(k);
 * int value = obj.addNum(num);
 */

Optimal Approach

In the optimal approach, we will use a max - heap. We can keep pushing elements in the heap till its size is less than k. When its size exceeds k, we pop the elements till the size equals k. Then we return the top element of the heap.

Analysis

• Time Complexity: O(logn)
• Space Complexity: O(n)

Implementation

C++

class KthLargest {
public:
    /** initialize your data structure here. */
	priority_queue<int, vector<int>, greater<int>> pq;
	int K;
    KthLargest(int k) {
		while(!pq.empty()) {
			pq.pop();
		}
		K = k;
    }
    
    int add(int num) { 
		pq.push(num);
		if(pq.size() < K) {
			return -1;
		}
        while(pq.size() > K) {
            pq.pop();
        }
		return pq.top();
    }
};

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest* obj = new KthLargest(k);
 * int value = obj->addNum(num);
 */

Java

class KthLargest {
    /** initialize your data structure here. */
	PriorityQueue<Integer> pq;
	int k;
    public KthLargest(int k) {
		this.k = k;
        pq = new PriorityQueue<>();
    }
    
    public int add(int num) {
		pq.add(num);
		if(pq.size() < k) {
			return -1;
		}
        while(pq.size() > k) {
            pq.poll();
        }
		return pq.peek();
    }
}

/**
 * Your KthLargest object will be instantiated and called as such:
 * KthLargest obj = new KthLargest(k);
 * int value = obj.addNum(num);
 */