View on NeetCode
View on LeetCode

Problem

Given an integer array nums and an integer k, return the kth largest element in the array.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

Can you solve it without sorting?

Example 1:

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

Example 2:

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

Constraints:

  • 1 <= k <= nums.length <= 10^5
  • -10^4 <= nums[i] <= 10^4

Solution

Approach 1: Min Heap of Size K

The key insight is to maintain a min heap of size k with the k largest elements. The root is the kth largest.

Implementation

import heapq

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # Min heap of size k
        heap = []

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

            # Keep only k largest
            if len(heap) > k:
                heapq.heappop(heap)

        return heap[0]

Approach 2: QuickSelect (Optimal Average)

import random

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        k = len(nums) - k  # Convert to kth smallest (0-indexed)

        def quickselect(left, right):
            pivot = random.randint(left, right)
            nums[pivot], nums[right] = nums[right], nums[pivot]

            # Partition
            store_index = left
            for i in range(left, right):
                if nums[i] < nums[right]:
                    nums[i], nums[store_index] = nums[store_index], nums[i]
                    store_index += 1

            nums[store_index], nums[right] = nums[right], nums[store_index]

            # Recurse on appropriate partition
            if store_index == k:
                return nums[store_index]
            elif store_index < k:
                return quickselect(store_index + 1, right)
            else:
                return quickselect(left, store_index - 1)

        return quickselect(0, len(nums) - 1)

Complexity Analysis

Min Heap:

  • Time Complexity: O(n log k), where n is array length. Each heap operation is O(log k).
  • Space Complexity: O(k) for the heap.

QuickSelect:

  • Time Complexity: O(n) average, O(n²) worst case. Average case partitions reduce problem size by half each time.
  • Space Complexity: O(1) with in-place partitioning.

Key Insights

  1. Kth Largest = (n-k)th Smallest: Converting to find the (n-k)th smallest simplifies QuickSelect implementation.

  2. Min Heap Strategy: A min heap of size k maintains the k largest elements, with the smallest (kth largest) at the root.

  3. QuickSelect Efficiency: QuickSelect is faster on average than heap for single queries, using partitioning like QuickSort.

  4. Randomization: Random pivot selection in QuickSelect avoids worst-case O(n²) on already sorted arrays.

  5. When to Use Each: Heap is better for multiple queries or when k is small. QuickSelect is better for single queries with large k.