Leetcode 215. Kth Largest Element in Array

randomized_quickselect.py

# This is the same thing as finding the (len(nums) - (k-1))th order statistic.
#   - Note: We need to transform "Kth Largest, 1-indexed" into "Kth Smallest, 0-indexed"
# Idea: Use randomized quickselect! It has O(n) expected time

import random


class Solution:

    # O(n) time, O(1) space
    def partition(self, nums: List[int], l: int, r: int) -> int:
        # Pick pivot at random, uniformly
        p_idx = random.randrange(l, r + 1)

        # Switch pivot into first position
        nums[l], nums[p_idx] = nums[p_idx], nums[l]

        # Maintain invariant that all elements before index i are <= pivot
        i, j = l, l + 1
        while j <= r:
            if nums[j] <= nums[l]:
                nums[j], nums[i+1] = nums[i+1], nums[j]
                i += 1
            j += 1

        # Switch pivot into proper place
        nums[l], nums[i] = nums[i], nums[l]
        return i

    # O(n) expected time
    def randomized_quickselect(self, nums: List[int], k_prime: int, l: int, r: int) -> int:
        if l == r:
            return nums[l]
        j = self.partition(nums, l, r)
        if j == k_prime:
            return nums[j]
        elif j < k_prime:
            l, r = j + 1, r
        else:
            l, r = l, j - 1
        return self.randomized_quickselect(nums, k_prime, l, r)

    def findKthLargest(self, nums: List[int], k: int) -> int:
        l, r, k_prime = 0, len(nums) - 1, len(nums) - k
        return self.randomized_quickselect(nums, k_prime, l, r)