Lomuto partition scheme in Python
Partition scheme that scans once and places the pivot at its final index; simple but does more swaps.
- File
lomuto_partition.py, 31 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/lomuto-partition/. The authors are listed in prov/N086ME.jsonl.
from typing import Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def lomuto_partition(a: list[T], lo: int = 0, hi: int | None = None) -> int:
"""Partition a[lo:hi] around x = a[hi - 1] and return the final index p of x.
After the call, each item of a[lo:p] is <= x, a[p] is x, and each item of a[p + 1:hi] is > x.
The range must not be empty.
"""
if hi is None:
hi = len(a)
if hi - lo < 1:
raise ValueError("the range is empty")
x = a[hi - 1] # @line L1
i = lo # @line L2
for j in range(lo, hi - 1): # @line L3
if not x < a[j]: # @line L4: a[j] <= x, with only <
a[i], a[j] = a[j], a[i] # @line L5
i += 1 # @line L6
a[i], a[hi - 1] = a[hi - 1], a[i] # @line L7
return i # @line L8
test_lomuto_partition.py 66 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/lomuto-partition/. The authors are listed in prov/N086ME.jsonl.
import itertools
import random
import pytest
from lomuto_partition import lomuto_partition
def check(a: list[int], lo: int, hi: int) -> None:
old = list(a)
x = a[hi - 1]
p = lomuto_partition(a, lo, hi)
assert lo <= p < hi
assert a[p] == x
assert all(v <= x for v in a[lo:p])
assert all(v > x for v in a[p + 1 : hi])
assert sorted(a[lo:hi]) == sorted(old[lo:hi])
assert a[:lo] == old[:lo] and a[hi:] == old[hi:]
def test_example() -> None:
a = [6, 3, 9, 1, 7, 2, 8, 5]
assert lomuto_partition(a) == 3
assert a == [3, 1, 2, 5, 7, 9, 8, 6]
def test_one_item() -> None:
a = [4]
assert lomuto_partition(a) == 0
assert a == [4]
def test_empty_range() -> None:
with pytest.raises(ValueError):
lomuto_partition([], 0, 0)
def test_equal_items_go_left() -> None:
a = [7, 7, 7, 7]
assert lomuto_partition(a) == 3
def test_extreme_pivots() -> None:
a = [1, 2, 3, 4, 5]
assert lomuto_partition(a) == 4
b = [5, 4, 3, 2, 1]
assert lomuto_partition(b) == 0
assert b == [1, 4, 3, 2, 5]
def test_all_small_inputs() -> None:
for n in range(1, 7):
for p in itertools.product(range(3), repeat=n):
check(list(p), 0, n)
def test_subrange() -> None:
rng = random.Random(11)
for _ in range(400):
a = [rng.randint(-20, 20) for _ in range(rng.randint(1, 40))]
lo = rng.randrange(len(a))
hi = rng.randint(lo + 1, len(a))
check(a, lo, hi)
What the next level needs
- the tests did not run in this build