A simplified quicksort implementation

quicksort.py

from typing import List, Tuple
import random


def quicksort(array: List[int]) -> List[int]:
    """Recursive implementation of quicksort.

    Args:
        array (List[int]): a list of integers

    Returns:
        List[int]: sorted array
    """

    # Base case
    if len(array) <= 1:
        return array

    # Step 1: choose a pivot
    pivot_idx = random.choice(range(len(array)))

    # Step 2: partition
    smaller_subarray, greater_subarray = partition(array, pivot_idx)

    # Step 3: Recurse on smaller and greater subarrays.
    return (
        quicksort(smaller_subarray)
        + [array[pivot_idx]]
        + quicksort(greater_subarray)
    )


def partition(array: List[int], pivot_idx: int) -> Tuple[List[int], List[int]]:
    """Parition array into subarrays smaller and greater than the pivot.

    Args:
        array (List[int]): input array
        pivot_idx (int): index of the pivot

    Returns:
        Tuple[List[int], List[int]]: smaller subarray, greater subarray
    """
    smaller_subarray, greater_subarray = [], []
    for idx, item in enumerate(array):
        # we don't want to add pivot to any of the sub-arrays
        if idx == pivot_idx:
            continue
        if item <= array[pivot_idx]:
            smaller_subarray.append(item)
        else:
            greater_subarray.append(item)
    return smaller_subarray, greater_subarray


if __name__ == "__main__":
    test_cases = [
        [],  # empty
        [1],  # single item
        [2, 1],  # 2 items
        [3, 2, 1],  # sorted in decreasing order
        [1, 1, 1],  # identical
        [random.randint(1, 1e6) for _ in range(100)],  # random integers
    ]

    for test_arr in test_cases:
        assert quicksort(test_arr) == sorted(test_arr)