Floyd-Rivest selection in Python
Selection that samples the input to choose two pivots bracketing the target rank, using close to n comparisons on average.
- File
floyd_rivest_select.py, 69 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/floyd-rivest-select/. The authors are listed in prov/4QBG0P.jsonl.
import math
import random
from typing import Optional, Protocol, Sequence, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def floyd_rivest_select(
a: Sequence[T], k: int, cutoff: int = 600, rng: Optional[random.Random] = None
) -> T:
"""Return the item of rank k (1-based) of a, the k-th smallest.
Lists of at most cutoff items are sorted. Larger lists are cut down by two
pivots u <= v chosen from a random sample so that the answer very likely
lies between them. The input is not changed. Uses only the operator <.
"""
n = len(a)
if not 1 <= k <= n:
raise ValueError("k must be between 1 and len(a)")
if rng is None:
rng = random.Random()
return _select(list(a), k, max(cutoff, 2), rng)
def _select(s: list[T], k: int, cutoff: int, rng: random.Random) -> T:
n = len(s)
if n <= cutoff: # @line L1
return sorted(s)[k - 1] # @line L1
m = min(n - 1, math.ceil(n ** (2 / 3))) # @line L2
sample = rng.sample(s, m) # @line L2
g = math.sqrt(m * math.log(n)) / 2 # @line L3
lo = max(1, math.floor(k * m / n - g)) # @line L4
hi = min(m, math.ceil(k * m / n + g)) # @line L4
u = _select(sample, lo, cutoff, rng) # @line L5
v = _select(sample, hi, cutoff, rng) # @line L5
small: list[T] = [] # @line L6
mid: list[T] = [] # @line L6
large: list[T] = [] # @line L6
if 2 * k <= n: # @line L7
for x in s: # @line L7
if v < x: # @line L7
large.append(x) # @line L7
elif x < u: # @line L7
small.append(x) # @line L7
else:
mid.append(x) # @line L7
else:
for x in s: # @line L8
if x < u: # @line L8
small.append(x) # @line L8
elif v < x: # @line L8
large.append(x) # @line L8
else:
mid.append(x) # @line L8
if k <= len(small): # @line L9
return _select(small, k, cutoff, rng) # @line L9
if k > len(small) + len(mid): # @line L10
return _select(large, k - len(small) - len(mid), cutoff, rng) # @line L10
if len(mid) == n: # @line L11
return sorted(s)[k - 1] # @line L11
return _select(mid, k - len(small), cutoff, rng) # @line L12
test_floyd_rivest_select.py 43 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/floyd-rivest-select/. The authors are listed in prov/4QBG0P.jsonl.
import random
import pytest
from floyd_rivest_select import floyd_rivest_select
def test_example() -> None:
a = [31, 41, 59, 26, 53, 58, 97, 93, 23, 84]
for k in range(1, 11):
assert floyd_rivest_select(a, k, cutoff=2, rng=random.Random(k)) == sorted(a)[k - 1]
assert a == [31, 41, 59, 26, 53, 58, 97, 93, 23, 84]
def test_edges() -> None:
assert floyd_rivest_select([5], 1) == 5
assert floyd_rivest_select([4] * 50, 17, cutoff=2) == 4
assert floyd_rivest_select([2, 1] * 40, 41, cutoff=2) == 2
with pytest.raises(ValueError):
floyd_rivest_select([1, 2], 0)
with pytest.raises(ValueError):
floyd_rivest_select([1, 2], 3)
def test_random_against_sorted() -> None:
rng = random.Random(1975)
for _ in range(1500):
n = rng.randint(1, 120)
a = [rng.randint(-30, 30) for _ in range(n)]
k = rng.randint(1, n)
cutoff = rng.choice([2, 3, 5, 10])
assert floyd_rivest_select(a, k, cutoff=cutoff, rng=rng) == sorted(a)[k - 1]
def test_large_default_cutoff() -> None:
rng = random.Random(489)
for _ in range(5):
a = [rng.random() for _ in range(20000)]
k = rng.randint(1, len(a))
assert floyd_rivest_select(a, k, rng=rng) == sorted(a)[k - 1]
What the next level needs
- the tests did not run in this build