Quickselect in Python
Partition around a pivot and recurse only into the side containing the target rank; expected O(n) time.
- File
quickselect.py, 50 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/quickselect/. The authors are listed in prov/TTGB6K.jsonl.
import random
from typing import Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def quickselect(a: list[T], k: int, rng: random.Random | None = None) -> T:
"""Return the k-th smallest item of a, for 1 <= k <= len(a).
The list is rearranged in place: after the call, a[k - 1] holds the answer,
the items before it are at most the answer and the items after it are at
least the answer. The pivot is uniform at random, so the expected number of
comparisons is less than 4n on n distinct items; the worst case is Theta(n^2).
"""
n = len(a)
if not 1 <= k <= n:
raise ValueError("k must satisfy 1 <= k <= len(a)")
rng = rng if rng is not None else random.Random()
lo, hi, t = 0, n, k - 1 # @line L1
while hi - lo > 1: # @line L2
r = rng.randrange(lo, hi) # @line L3
a[r], a[hi - 1] = a[hi - 1], a[r] # @line L4
p = partition(a, lo, hi) # @line L5
if t == p: # @line L6
return a[p]
if t < p: # @line L7
hi = p
else:
lo = p + 1 # @line L8
return a[lo] # @line L9
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 of x."""
x = a[hi - 1] # @line L10
i = lo # @line L11
for j in range(lo, hi - 1): # @line L12
if not x < a[j]: # @line L13: a[j] <= x
a[i], a[j] = a[j], a[i] # @line L14
i += 1 # @line L15
a[i], a[hi - 1] = a[hi - 1], a[i] # @line L16
return i # @line L17
test_quickselect.py 52 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/quickselect/. The authors are listed in prov/TTGB6K.jsonl.
import random
import pytest
from quickselect import partition, quickselect
def test_example() -> None:
a = [7, 2, 9, 4, 1, 8, 3, 6, 5]
assert quickselect(list(a), 5, random.Random(1)) == 5
for k in range(1, 10):
assert quickselect(list(a), k, random.Random(k)) == k
def test_partition() -> None:
a = [7, 2, 9, 4, 1, 8, 3, 6, 5]
p = partition(a, 0, len(a))
assert p == 4 and a[p] == 5
assert all(x <= 5 for x in a[:p]) and all(x > 5 for x in a[p + 1 :])
def test_edges() -> None:
assert quickselect([42], 1) == 42
assert quickselect([2, 1], 1) == 1
assert quickselect([2, 1], 2) == 2
assert quickselect([3, 3, 3, 3], 2) == 3
with pytest.raises(ValueError):
quickselect([], 1)
with pytest.raises(ValueError):
quickselect([1, 2], 3)
def test_rearranges_around_answer() -> None:
rng = random.Random(7)
a = [rng.randint(0, 9) for _ in range(40)]
v = quickselect(a, 17, rng)
assert a[16] == v
assert all(x <= v for x in a[:16]) and all(x >= v for x in a[17:])
def test_random_against_sorted() -> None:
rng = random.Random(9150)
for _ in range(800):
n = rng.randint(1, 50)
a = [rng.randint(-20, 20) for _ in range(n)]
k = rng.randint(1, n)
b = list(a)
assert quickselect(b, k, rng) == sorted(a)[k - 1]
assert sorted(b) == sorted(a)
What the next level needs
- the tests did not run in this build