Counting sort in Python
Sorts n integers in [0,k) in O(n+k) time by counting keys and prefix-summing the counts.
- File
counting_sort.py, 33 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/counting-sort/. The authors are listed in prov/FQMV6V.jsonl.
from typing import Callable, TypeVar
T = TypeVar("T")
def _identity(x: int) -> int:
return x
def counting_sort(a: list[T], k: int, key: Callable[[T], int] = _identity) -> list[T]: # type: ignore[assignment]
"""Return a new list with the items of a, stably sorted by key.
Precondition: 0 <= key(x) < k for each item x. The cost is
Theta(n + k) time and Theta(n + k) extra space for n = len(a).
"""
n = len(a)
c = [0] * k # @line L1
for i in range(n): # @line L2
c[key(a[i])] += 1 # @line L3
s = 0 # @line L4
for v in range(k): # @line L5
t = c[v] # @line L6
c[v] = s
s += t
b: list[T | None] = [None] * n # @line L7
for i in range(n): # @line L8
v = key(a[i])
b[c[v]] = a[i] # @line L9
c[v] += 1 # @line L10
return b # type: ignore[return-value] # @line L11
test_counting_sort.py 39 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/counting-sort/. The authors are listed in prov/FQMV6V.jsonl.
import random
from counting_sort import counting_sort
def test_example() -> None:
assert counting_sort([2, 5, 3, 0, 2, 3, 0, 3], 6) == [0, 0, 2, 2, 3, 3, 3, 5]
def test_edges() -> None:
assert counting_sort([], 1) == []
assert counting_sort([0], 1) == [0]
assert counting_sort([4, 4, 4], 5) == [4, 4, 4]
assert counting_sort([3, 2, 1, 0], 4) == [0, 1, 2, 3]
assert counting_sort([0, 9], 10) == [0, 9]
def test_stable() -> None:
a = [(2, "a"), (1, "b"), (2, "c"), (0, "d"), (1, "e"), (2, "f")]
out = counting_sort(a, 3, key=lambda p: p[0])
assert out == sorted(a, key=lambda p: p[0])
assert [t for _, t in out] == ["d", "b", "e", "a", "c", "f"]
def test_input_unchanged() -> None:
a = [3, 1, 2]
counting_sort(a, 4)
assert a == [3, 1, 2]
def test_random_against_sorted() -> None:
rng = random.Random(8021)
for _ in range(600):
k = rng.randint(1, 20)
a = [(rng.randrange(k), j) for j in range(rng.randint(0, 50))]
assert counting_sort(a, k, key=lambda p: p[0]) == sorted(a, key=lambda p: p[0])
What the next level needs
- the tests did not run in this build