Problem Statement in English

Given an array nums you have to return the $k\text{th}$ largest element in it.

Approach

If we use a min heap, and only maintain the $k\text{th}$ largest elements in it, then the first element in the min heap will be our answer.

For example, if we have nums: [3,2,1,5,6,4], and k = $2$, and we trace it out:

  • minHeap = [3]
  • minHeap = [2,3]
  • minHeap = [1,2,3], but since len(minHeap) = 3, and we only need to track the 2 largest elements, we can pop the first element from the heap, so minHeap = [2,3].
  • minHeap = [2,3,5], but since we only need the $2$ largest, we pop the first element, so minHeap = [3,5].

and so on.

Maintaining only $k$ elements is important because we can see that according to the constraints we can have upto $10^5$ elements in the array, and there’s no reason to store so many in an array again. That’s an expensive waste of space.

Solution in Python


class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        minHeap = []

        for num in nums:
            heapq.heappush(minHeap, num)

            if len(minHeap) > k:
                heapq.heappop(minHeap)

        return minHeap[0]

Note: Python heapqs are min heaps.

Complexity

  • Time: $O(nlog(k))$
    Heap push takes $log(k)$ time ($k$ because we only have $k$ elements in the array at any time), and for pushing $n$ such elements, that’s $nlog(k)$.

  • Space: $O(k)$
    Since we only store $k$ elements in the heap.

And we are done.