Kth Largest Element in an Array
- 3 min read
- LC-Medium
- Binary-heap
Solution
- About the best we can due in worst case time complexity is O(n + k log n). See less optimal solution below
- There is a solution with a better average case time complexity of O(n) but worst case O(n^2)
- The algorithm is called Quick Select, pretty similar to Quick Sort
- The solution TLEs on Leetcode but it is faster on average case
- Basically, at the end of each partition operation in quick sort / quick select, the partition index alone is in the correct position
- We initialize
n = len(nums) - k. The element at index n is the kth largest element in the array - Therefore, until our partition index is equal to n, we keep recurring on the subarrays after partitioning depending on the partition index
- If the value at partition index is less than n, we need to recur to the right of the array. Vice versa for the left
- If you're confused about the partitioning process, look at the mycodeschool video for quicksort.
- Basically, the partitioning is to make it so that all elements lesser than the pivot are to the left of the array and all elements greater are to the right
- The initialization and changing of the
i and p pointers are what are most confusing, so look at the mycodeschool video ok.
| Time | Space | Explanation |
|---|
O(n^2) worst case, O(n) average case | O() | |
def findKthLargest(self, nums: List[int], k: int) -> int:
n = len(nums) - k
def quickSelect(l, r):
p = l
for i in range(l, r):
if nums[i] <= nums[r]:
nums[i], nums[p] = nums[p], nums[i]
p += 1
nums[p], nums[r] = nums[r], nums[p]
if p < n:
return quickSelect(p + 1, r)
elif p > n:
return quickSelect(l, p - 1)
return nums[p]
return quickSelect(0, len(nums) - 1)
Less Optimal Solution
- Similar to Kth Largest Element in a Stream, heapify the input array (min heap), pop until size of heap is k, and then return the first element in the heap
| Time | Space | Explanation |
|---|
O(n + k log n) | O() | O(n) to build the heap, log (n) for each heappop. and we heappop k times |
def findKthLargest(self, nums: List[int], k: int) -> int:
heapify(nums)
while len(nums) > k:
heappop(nums)
return nums[0]