Leetcode 912. Sort an Array (Randomized QuickSort)

randomized_quicksort.py

# Let's implement Randomized QuickSort!
import random

class Solution:

    # O(n)
    def partition(self, nums: List[int], l: int, r: int) -> int:
        # Choose random pivot, put it first in the array 
        random_idx = random.choice(range(l, r + 1))
        nums[l], nums[random_idx] = nums[random_idx], nums[l]
        pivot = l

        # Maintain loop invariant: for all non-pivot elements, everything
        #   1. before i is is < nums[pivot] 
        #   2. between i, j is > nums[pivot]
        #   3. after j is yet to be processed 
        i, j = l, l + 1
        while j <= r:
            if nums[j] < nums[pivot]:
                nums[i + 1], nums[j] = nums[j], nums[i + 1]
                i += 1
            j += 1
        
        # Switch pivot into correct place
        nums[i], nums[pivot] = nums[pivot], nums[i]
        return i

    # O(n log n) expected time, O(log n) space
    def quicksort(self, nums: List[int], l: int, r: int) -> None:
        if 0 <= l < r < len(nums):
            p = self.partition(nums, l, r)
            self.quicksort(nums, l, p - 1)
            self.quicksort(nums, p + 1, r)

    def sortArray(self, nums: List[int]) -> List[int]:
        self.quicksort(nums, 0, len(nums)-1)
        return nums

Note: Leetcode will make sure to test the worst cases that make QuickSort run in $O(n^2)$ time! So this will TLE, for example, when all of the large number of elements are uniformly the same.