Kth Largest Element

Finding the Kth Largest Element Using a Min Heap

The problem of finding the Kth largest element in an unsorted list is a common interview question. A common and efficient approach to solve this problem is using a min heap of size K.

Concept

The key idea is that the Kth largest element is the smallest element in a list of the K largest elements of the array.

Example

Given the array nums = [3, 2, 1, 5, 6, 4] and k = 2, we need to find the 2nd largest element. The sorted array would be [1, 2, 3, 4, 5, 6], and the 2nd largest element is 5.

Now, consider maintaining a list of the two largest elements:

  • [5, 6]

Here, the 2nd largest element is the smallest element of [5, 6], which is 5.

Approach

  1. Min Heap: Use a min heap to store the K largest elements encountered during the iteration.
  2. Heap Operations:
    • As you iterate through the list, push elements into the heap.
    • If the heap size exceeds K, remove the smallest element (this maintains the K largest elements).
  3. Final Result: After processing all elements, the smallest element in the heap will be the Kth largest element.

Code Explanation

import heapq
from typing import List

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # Min heap to store the K largest elements
        heap = []

        for num in nums:
            # Push the current number into the heap if heap has less than K elements
            if len(heap) < k:
                heapq.heappush(heap, num)
                continue

            # If the current number is larger than the smallest in the heap, replace it
            if heap[0] < num:
                heapq.heappop(heap)  # Remove the smallest element
                heapq.heappush(heap, num)  # Push the current number

        # The root of the heap (smallest element in the heap) is the Kth largest element
        return heap[0]

Explanation of Steps:

  1. Initialization: We start with an empty heap.
  2. Heap Insertion: For each element in nums, if the heap has fewer than k elements, we add the element to the heap.
  3. Replacement: If the heap already has k elements, we compare the current element with the smallest element in the heap (root). If the current element is larger, we remove the root and add the current element to maintain the heap size at k.
  4. Result: After all elements are processed, the root of the heap (the smallest element) is the Kth largest element.

Time Complexity

  • Heap Operations: Inserting and removing elements from the heap takes O(log K) time.
  • Overall Complexity: The overall time complexity is O(N log K), where N is the number of elements in the array.

This approach efficiently finds the Kth largest element with minimal memory overhead and optimal performance.