Kth Largest Element in Array

Min-Heap

Problem

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

For example:

Input: [3,2,1,5,6,4] and k = 2
Output: 5

Input: [3,2,3,1,2,4,5,5,6] and k = 4
Output: 4

Thought Process

• With a min-heap, the top element (root) will be the Kth largest if we do the respective comparisons since the larger numbers will always be below it (as children); with a max-heap, the top element (root) will be the Kth smallest if we do the respective comparisons since the smaller numbers will be below it

• A min-heap is originally storing the smallest elements, with the absoute smallest at top. We can discard this smallest element in place of a new one if we do comparisons with the other elements in the array to see if the other elements are larger.

Solution

from heapq import *
class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = []
        
        for i in range(k):
            heappush(heap, nums[i])
            
        for i in range(k, len(nums)):
            if nums[i] > heap[0]:
                heappop(heap)
                heappush(heap, nums[i])
        return heap[0]
        

Similar Question: Find Kth Smallest Element

from heapq import *


def find_Kth_smallest_number(nums, k):
  maxHeap = []
 
  for i in range(k):
    heappush(maxHeap, -nums[i])
  
  for i in range(k, len(nums)):
    if -nums[i] > maxHeap[0]:
      heappop(maxHeap)
      heappush(maxHeap, -nums[i])

  return -maxHeap[0]

• We use a max-heap to find the Kth smallest element because as we know, the root is the biggest element in the max heap. So, since we want to keep track of the ‘K’ smallest numbers, we can compare every number with the root while iterating through all numbers, and if it is smaller than the root, we’ll take the root out and insert the smaller number.

Key Points

• A heap will always store the largest/smallest elements with the Kth smallest or largest element being at the top of the heap (as the root).

Time Complexity

• Time: $O(N*log(K))$

• Space: $O(K)$