3-way partition quickselect quicksort kth smallest

Find the kth largest element in an unsorted array.

Examples

[3,2,1,5,6,4] k=2 -> 5

[3,2,3,1,2,4,5,5,6] k=4 -> 4

3-Way Partition

3-Way Partition (see Dutch National Flag problem) is my preferred partition algorithm.

  2   4   6   3   4   7   8   3   pivot=3
 l,m                          r
  2   4   6   3   4   7   8   3
     l,m                      r
  2   3   6   3   4   7   8   4
     l,m                  r
  2   3   6   3   4   7   8   4
      l   m               r
  2   3   8   3   4   7   6   4
      l   m           r
  2   3   7   3   4   8   6   4
      l   m       r
  2   3   4   3   7   8   6   4
      l   m   r
  2   3   3   4   7   8   6   4
      l  m,r
  2   3   3   4   7   8   6   4
      l   r   m
 l-1         r+1

def partition(nums, l, r):
    pivot = nums[r]
    mid = l
    while mid <= r:
        if nums[mid] < pivot:
            nums[l], nums[mid] = nums[mid], nums[l]
            l += 1
            mid += 1
        elif nums[mid] > pivot:
            nums[r], nums[mid] = nums[mid], nums[r]
            r -= 1
        else:
            mid += 1
    return l-1, r+1

Solution

Quickselect is a selection algorithm to find the kth smallest element in an unsorted array.

def findKthLargest(nums, k):
    def helper(nums, l, r, k): # kth smallest
        a, b = partition(nums, l, r)
        if a < k-1 and k-1 < b: # a < k-1 < b
            return nums[k-1]
        elif a >= k-1:
            return helper(nums, l, a, k)
        else:
            return helper(nums, b, r, k)

    # kth largest equals to `len-k+1`th smallest
    return helper(nums, 0, len(nums)-1, len(nums)-k+1)

Additional Topic

Quick Sort

# best avg O(nlogn)
def quicksort(arr, l, r): 
    if l >= r or l < 0 or r >= len(arr):
        return
    a, b = partition(arr, l, r)
    quicksort(arr, l, a)
    quicksort(arr, b, r)