67. Kth Largest Element in an Array
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
-
Kth Largest = (n-k)th Smallest: Converting to find the (n-k)th smallest simplifies QuickSelect implementation.
-
Min Heap Strategy: A min heap of size k maintains the k largest elements, with the smallest (kth largest) at the root.
-
QuickSelect Efficiency: QuickSelect is faster on average than heap for single queries, using partitioning like QuickSort.
-
Randomization: Random pivot selection in QuickSelect avoids worst-case O(n²) on already sorted arrays.
-
When to Use Each: Heap is better for multiple queries or when k is small. QuickSelect is better for single queries with large k.