Kth Largest Element in an Array

Kth Largest Element in an Array

January 4, 2024
medium
Quickselect

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.

Example: # 

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

Solution #

Algorithm: Quickselect

This approach is based on the QuickSort algorithm. In QuickSort, we pick a pivot and partition the array around the pivot. The idea of Quickselect is similar, but instead of recursively doing both sides of the pivot, as in QuickSort, Quickselect only recurses into one side – the side with the element it is searching for. This reduces the average complexity from O(n log n) to O(n), with a worst-case of O(n^2).

Algorithm Steps: #

  1. Choose a Pivot: Randomly select a pivot element from the array.
  2. Partition: Reorder the array so that all elements with values less than the pivot come before the pivot, while all elements with values greater than the pivot come after it.
  3. Target Index Identification:
    • If the pivot’s position is k-1, we found the kth largest element.
    • If the pivot’s position is greater than k-1, repeat the process in the left subarray.
    • If the pivot’s position is less than k-1, repeat the process in the right subarray.

Python Code # 

import random

def findKthLargest(nums, k):
    def partition(left, right, pivot_index):
        pivot = nums[pivot_index]
        # 1. move pivot to end
        nums[pivot_index], nums[right] = nums[right], nums[pivot_index]  
        
        # 2. move all smaller elements to the left
        store_index = left
        for i in range(left, right):
            if nums[i] < pivot:
                nums[store_index], nums[i] = nums[i], nums[store_index]
                store_index += 1

        # 3. move pivot to its final place
        nums[right], nums[store_index] = nums[store_index], nums[right]  
        
        return store_index

    def select(left, right, k_smallest):
        """
        Returns the k-th smallest element of list within left..right
        """
        if left == right:       # If the list contains only one element,
            return nums[left]   # return that element

        # select a random pivot_index
        pivot_index = random.randint(left, right)     
                            
        # find the pivot position in a sorted list   
        pivot_index = partition(left, right, pivot_index)
        
        # the pivot is in its final sorted position
        if k_smallest == pivot_index:
             return nums[k_smallest]
        # go left
        elif k_smallest < pivot_index:
            return select(left, pivot_index - 1, k_smallest)
        # go right
        else:
            return select(pivot_index + 1, right, k_smallest)

    return select(0, len(nums) - 1, len(nums) - k)

# Example Usage
nums = [3,2,1,5,6,4]
k = 2
print(findKthLargest(nums, k))
# Output: 5

Time Complexity: #

  • Average case: O(n)
  • Worst case: O(n^2)

Space Complexity: #

  • O(1): The space complexity is constant as the partitioning is done in-place.