Randomized quicksort in Python

Quicksort with uniformly random pivots, giving expected O(n log n) on every input.

File: randomized_quicksort.py, 50 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

randomized_quicksort.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/randomized-quicksort/. The authors are listed in prov/PP7XPG.jsonl.

import random
import sys
from typing import Protocol, TypeVar


class Ordered(Protocol):
    def __lt__(self, other: "Ordered", /) -> bool: ...


T = TypeVar("T", bound=Ordered)

# On an array of equal items the Lomuto partition is one-sided, and the
# recursion is n - 1 calls deep whatever the pivots are. Python stops at 1000
# calls by default, so this module raises the limit.
sys.setrecursionlimit(max(sys.getrecursionlimit(), 20000))


def randomized_quicksort(a: list[T], rng: random.Random | None = None) -> None:
    """Sort a in place, in non-decreasing order, with uniformly random pivots.

    Pass a seeded random.Random as rng to make a run repeatable.
    """
    if rng is None:
        rng = random.Random()
    _sort(a, 0, len(a), rng)


def _sort(a: list[T], lo: int, hi: int, rng: random.Random) -> None:
    if hi - lo <= 1:  # @line L1
        return
    r = rng.randrange(lo, hi)  # @line L2
    a[r], a[hi - 1] = a[hi - 1], a[r]  # @line L3
    p = partition(a, lo, hi)  # @line L4
    _sort(a, lo, p, rng)  # @line L5
    _sort(a, p + 1, hi, rng)  # @line L6


def partition(a: list[T], lo: int, hi: int) -> int:
    """Lomuto partition of a[lo:hi] around x = a[hi - 1]. Return the final index p of x."""
    x = a[hi - 1]
    i = lo
    for j in range(lo, hi - 1):
        if not x < a[j]:  # a[j] <= x, with only <
            a[i], a[j] = a[j], a[i]
            i += 1
    a[i], a[hi - 1] = a[hi - 1], a[i]
    return i

test_randomized_quicksort.py 71 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/randomized-quicksort/. The authors are listed in prov/PP7XPG.jsonl.

import itertools
import random

from randomized_quicksort import randomized_quicksort


class Counted:
    count = 0

    def __init__(self, v: int) -> None:
        self.v = v

    def __lt__(self, other: "Counted") -> bool:
        Counted.count += 1
        return self.v < other.v


def test_example() -> None:
    a = [1, 2, 3, 4, 5, 6, 7, 8]
    randomized_quicksort(a, random.Random(1))
    assert a == [1, 2, 3, 4, 5, 6, 7, 8]


def test_edges() -> None:
    for a in ([], [1], [2, 1], [1, 1, 1], [3, -1, 3, 0]):
        b = list(a)
        randomized_quicksort(b, random.Random(0))
        assert b == sorted(a)


def test_all_permutations() -> None:
    rng = random.Random(3)
    for n in range(7):
        for p in itertools.permutations(range(n)):
            a = list(p)
            randomized_quicksort(a, rng)
            assert a == list(range(n))


def test_random_with_duplicates() -> None:
    rng = random.Random(11)
    for _ in range(200):
        a = [rng.randint(0, 20) for _ in range(rng.randint(0, 60))]
        b = list(a)
        randomized_quicksort(b, rng)
        assert b == sorted(a)


def test_default_rng() -> None:
    a = list(range(500, 0, -1))
    randomized_quicksort(a)
    assert a == list(range(1, 501))


def test_mean_on_sorted_input() -> None:
    # The expected count on distinct keys is 2(n + 1)H_n - 4n on every input.
    n = 200
    h = sum(1 / k for k in range(1, n + 1))
    expected = 2 * (n + 1) * h - 4 * n
    rng = random.Random(5)
    total = 0
    runs = 200
    for _ in range(runs):
        a = [Counted(i) for i in range(n)]
        Counted.count = 0
        randomized_quicksort(a, rng)
        total += Counted.count
    assert abs(total / runs - expected) < 0.05 * expected

What the next level needs

the tests did not run in this build

Read the full page about Randomized quicksort