Merge path in Python
Splits a merge into independent pieces by binary searching along the merge path of the two inputs.
- File
merge_path.py, 58 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/merge-path/. The authors are listed in prov/8R73CV.jsonl.
"""Merge path: split a stable merge into p pieces of equal output length.
The p pieces are independent; the code merges them one after the other,
which gives the same output as merging them at the same time.
"""
from typing import Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def co_rank(k: int, a: list[T], b: list[T]) -> tuple[int, int]:
"""Return (i, j) with i + j = k such that the first k items of the
stable merge of a and b (a first on ties) are a[:i] and b[:j]."""
m, n = len(a), len(b)
lo, hi = max(0, k - n), min(k, m) # @line L1
while lo < hi: # @line L2
i = (lo + hi) // 2 # @line L3
if b[k - i - 1] < a[i]: # @line L4
hi = i
else:
lo = i + 1
return lo, k - lo # @line L5
def sequential_merge(a: list[T], b: list[T]) -> list[T]:
"""The stable two-finger merge, a first on ties."""
out: list[T] = []
i = j = 0
while i < len(a) and j < len(b):
if b[j] < a[i]:
out.append(b[j])
j += 1
else:
out.append(a[i])
i += 1
return out + a[i:] + b[j:]
def merge_path(a: list[T], b: list[T], p: int) -> list[T]:
"""Merge sorted a and b stably with p independent pieces."""
if p < 1:
raise ValueError("need p >= 1")
total = len(a) + len(b)
out: list[T] = []
for t in range(p): # @line L6
i0, j0 = co_rank(t * total // p, a, b) # @line L7
i1, j1 = co_rank((t + 1) * total // p, a, b)
out.extend(sequential_merge(a[i0:i1], b[j0:j1])) # @line L8
return out
test_merge_path.py 64 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/merge-path/. The authors are listed in prov/8R73CV.jsonl.
import random
import pytest
from merge_path import co_rank, merge_path, sequential_merge
A = [1, 4, 6, 9]
B = [2, 3, 7, 8, 10]
def test_example_co_ranks() -> None:
assert [co_rank(k, A, B) for k in (0, 3, 6, 9)] == [(0, 0), (1, 2), (3, 3), (4, 5)]
def test_example() -> None:
assert merge_path(A, B, 3) == [1, 2, 3, 4, 6, 7, 8, 9, 10]
def test_edges() -> None:
assert merge_path([], [], 4) == []
assert merge_path([1, 2], [], 3) == [1, 2]
assert merge_path([], [1, 2], 5) == [1, 2]
with pytest.raises(ValueError):
merge_path([1], [2], 0)
def test_stable_ties() -> None:
a = [(2, "a0"), (5, "a1")]
b = [(2, "b0"), (5, "b1")]
class K:
def __init__(self, pair: tuple[int, str]) -> None:
self.k, self.tag = pair
def __lt__(self, other: "K") -> bool:
return self.k < other.k
ka, kb = [K(x) for x in a], [K(x) for x in b]
for p in (1, 2, 3, 4):
assert [x.tag for x in merge_path(ka, kb, p)] == ["a0", "b0", "a1", "b1"]
def test_co_rank_matches_merge_prefix() -> None:
rng = random.Random(1)
for _ in range(300):
a = sorted(rng.randint(0, 9) for _ in range(rng.randint(0, 12)))
b = sorted(rng.randint(0, 9) for _ in range(rng.randint(0, 12)))
tagged = sequential_merge([(x, 0) for x in a], [(x, 1) for x in b])
for k in range(len(a) + len(b) + 1):
i, j = co_rank(k, a, b)
assert i == sum(1 for _, src in tagged[:k] if src == 0)
assert i + j == k
def test_random_against_sorted() -> None:
rng = random.Random(2)
for _ in range(300):
a = sorted(rng.randint(-20, 20) for _ in range(rng.randint(0, 40)))
b = sorted(rng.randint(-20, 20) for _ in range(rng.randint(0, 40)))
assert merge_path(a, b, rng.randint(1, 10)) == sorted(a + b)
What the next level needs
- the tests did not run in this build