Balanced multiway merge sort in Python
Merge runs in passes over a fixed set of tapes or files, with the same number of inputs and outputs in each pass.
- File
balanced_merge_sort.py, 74 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/balanced-merge-sort/. The authors are listed in prov/K6SD9E.jsonl.
import heapq
from typing import Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
class _Head:
"""The head of the run from tape d, at position j. Equal keys order by tape number."""
__slots__ = ("x", "d", "j")
def __init__(self, x: T, d: int, j: int) -> None:
self.x, self.d, self.j = x, d, j
def __lt__(self, other: "_Head") -> bool:
if self.x < other.x:
return True
if other.x < self.x:
return False
return self.d < other.d
def _merge(runs: list[list[T]]) -> list[T]:
"""A stable merge of runs; on equal keys the run with the smaller index wins."""
heap = [_Head(r[0], d, 0) for d, r in enumerate(runs) if r]
heapq.heapify(heap)
out: list[T] = []
while heap:
h = heap[0]
out.append(h.x)
if h.j + 1 < len(runs[h.d]):
heapq.heapreplace(heap, _Head(runs[h.d][h.j + 1], h.d, h.j + 1))
else:
heapq.heappop(heap)
return out
def balanced_merge_sort(f: list[T], m: int = 4, t: int = 4) -> int:
"""Sort f in place with a balanced merge on t simulated tapes and a memory of m items.
A tape is a list of runs. Precondition: m >= 1 and t >= 4, so that the
merge order p = t // 2 is at least 2. The sort is stable. Returns the
number of merge passes, which is ceil(log_p s) for s = ceil(n / m) initial
runs (0 if s <= 1). Each pass reads and writes all n items; the time in
memory is O(n log n).
"""
if m < 1 or t < 4:
raise ValueError("need m >= 1 and t >= 4")
p = t // 2 # @line L1
a: list[list[list[T]]] = [[] for _ in range(p)]
j = 0 # @line L2
for start in range(0, len(f), m): # @line L3
a[j % p].append(sorted(f[start:start + m])) # @line L4
j += 1
passes = 0
while sum(len(tape) for tape in a) > 1: # @line L5
b: list[list[list[T]]] = [[] for _ in range(p)]
i = 0 # @line L6
while any(i < len(tape) for tape in a): # @line L7
group = [tape[i] for tape in a if i < len(tape)]
b[i % p].append(_merge(group)) # @line L8
i += 1
a = b # @line L9
passes += 1
f[:] = a[0][0] if a[0] else [] # @line L10
return passes
test_balanced_merge_sort.py 67 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/balanced-merge-sort/. The authors are listed in prov/K6SD9E.jsonl.
import random
import pytest
from balanced_merge_sort import balanced_merge_sort
def ceil_log(s: int, p: int) -> int:
k, q = 0, 1
while q < s:
q *= p
k += 1
return k
def test_example() -> None:
f = [12, 3, 8, 1, 14, 6, 2, 11, 9, 5, 13, 4, 10, 7]
assert balanced_merge_sort(f, 2, 4) == 3
assert f == list(range(1, 15))
def test_edges() -> None:
a: list[int] = []
assert balanced_merge_sort(a) == 0 and a == []
a = [1]
assert balanced_merge_sort(a) == 0 and a == [1]
a = [4, 3, 2, 1]
assert balanced_merge_sort(a, 4, 4) == 0 and a == [1, 2, 3, 4]
a = [5] * 9
balanced_merge_sort(a, 2, 6)
assert a == [5] * 9
a = list(range(30, 0, -1))
assert balanced_merge_sort(a, 1, 8) == 3 and a == list(range(1, 31))
with pytest.raises(ValueError):
balanced_merge_sort([1], 2, 3)
def test_stable() -> None:
class Key:
def __init__(self, k: int, tag: int) -> None:
self.k, self.tag = k, tag
def __lt__(self, other: "Key") -> bool:
return self.k < other.k
rng = random.Random(3)
for _ in range(50):
a = [Key(rng.randint(0, 3), i) for i in range(rng.randint(0, 60))]
want = sorted((x.k, x.tag) for x in a)
balanced_merge_sort(a, rng.randint(1, 5), rng.randint(4, 9))
assert [(x.k, x.tag) for x in a] == want
def test_random_against_sorted() -> None:
rng = random.Random(1961)
for _ in range(600):
n = rng.randint(0, 100)
m = rng.randint(1, 6)
t = rng.randint(4, 10)
f = [rng.randint(-40, 40) for _ in range(n)]
want = sorted(f)
passes = balanced_merge_sort(f, m, t)
assert f == want
assert passes == ceil_log(-(-n // m), t // 2)
What the next level needs
- the tests did not run in this build