Shiloach-Vishkin merging in Python
Merges two sorted lists with p processors in O(n/p + log n) PRAM time by ranking evenly spaced cuts of both lists in the other list.
- File
shiloach_vishkin_merge.py, 62 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/shiloach-vishkin-merge/. The authors are listed in prov/T0XRJS.jsonl.
"""Shiloach-Vishkin merging with p processors, simulated one processor at a time.
Both inputs are cut into blocks of ceil(m/p) and ceil(n/p) items. Each cut
is ranked in the other input by binary search. The cuts are put in output
order by arithmetic, and the pieces between them are merged one by one.
The merge is stable: on equal keys, items of A come first.
"""
import bisect
from typing import Protocol, TypeVar
class Ordered(Protocol):
def __lt__(self, other: "Ordered", /) -> bool: ...
T = TypeVar("T", bound=Ordered)
def ceil_div(x: int, y: int) -> int:
return -(-x // y)
def sv_splits(a: list[T], b: list[T], p: int) -> list[tuple[int, int]]:
"""Return the sorted splits (i, j): the first i + j output items are a[:i] and b[:j]."""
m, n = len(a), len(b)
sa, sb = max(1, ceil_div(m, p)), max(1, ceil_div(n, p)) # @line L1
a_cuts = [(t * sa, bisect.bisect_left(b, a[t * sa])) for t in range(1, p) if t * sa < m] # @line L2
b_cuts = [(bisect.bisect_right(a, b[t * sb]), t * sb) for t in range(1, p) if t * sb < n] # @line L3
order: list[tuple[int, int] | None] = [None] * (len(a_cuts) + len(b_cuts))
for idx, (i, j) in enumerate(a_cuts):
order[idx + max(0, ceil_div(j, sb) - 1)] = (i, j) # @line L4
for idx, (i, j) in enumerate(b_cuts):
order[idx + max(0, ceil_div(i, sa) - 1)] = (i, j)
return [(0, 0)] + order + [(m, n)] # type: ignore[operator] # @line L5
def merge_piece(a: list[T], b: list[T], out: list, i: int, i2: int, j: int, j2: int) -> None:
"""Sequential stable merge of a[i:i2] and b[j:j2] into out[i + j : i2 + j2]."""
k = i + j
while i < i2 or j < j2:
if j >= j2 or (i < i2 and not b[j] < a[i]):
out[k] = a[i]
i += 1
else:
out[k] = b[j]
j += 1
k += 1
def shiloach_vishkin_merge(a: list[T], b: list[T], p: int) -> tuple[list[T], list[tuple[int, int]]]:
"""Return the stable merge of the sorted lists a and b, and the splits used."""
if p < 1:
raise ValueError("need p >= 1")
s = sv_splits(a, b, p)
out: list = [None] * (len(a) + len(b))
for (i, j), (i2, j2) in zip(s, s[1:]): # @line L6
merge_piece(a, b, out, i, i2, j, j2) # @line L7
return out, s
test_shiloach_vishkin_merge.py 69 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/shiloach-vishkin-merge/. The authors are listed in prov/T0XRJS.jsonl.
import itertools
import math
import random
import pytest
from shiloach_vishkin_merge import shiloach_vishkin_merge
def stable_merge(a, b):
return sorted([(x, 0, i) for i, x in enumerate(a)] + [(x, 1, i) for i, x in enumerate(b)])
def test_example() -> None:
out, s = shiloach_vishkin_merge([1, 4, 6, 9], [2, 3, 7, 8, 10], 3)
assert out == [1, 2, 3, 4, 6, 7, 8, 9, 10]
assert s == [(0, 0), (2, 2), (3, 2), (4, 4), (4, 5)]
def test_edges() -> None:
assert shiloach_vishkin_merge([], [], 4)[0] == []
assert shiloach_vishkin_merge([1, 2], [], 3)[0] == [1, 2]
assert shiloach_vishkin_merge([], [1, 2], 3)[0] == [1, 2]
with pytest.raises(ValueError):
shiloach_vishkin_merge([1], [2], 0)
def test_small_exhaustive() -> None:
for m in range(5):
for n in range(5):
for a in itertools.product(range(3), repeat=m):
for b in itertools.product(range(3), repeat=n):
a2, b2 = sorted(a), sorted(b)
for p in (1, 2, 3, 5):
tagged_a = [(x, 0, i) for i, x in enumerate(a2)]
tagged_b = [(x, 1, i) for i, x in enumerate(b2)]
out, _ = shiloach_vishkin_merge(tagged_a, tagged_b, p)
assert out == stable_merge(a2, b2)
def test_piece_sizes() -> None:
rng = random.Random(7)
for _ in range(200):
m, n, p = rng.randrange(60), rng.randrange(60), rng.randrange(1, 12)
a = sorted(rng.randrange(20) for _ in range(m))
b = sorted(rng.randrange(20) for _ in range(n))
out, s = shiloach_vishkin_merge(a, b, p)
assert out == sorted(a + b)
assert len(s) - 1 <= 2 * p - 1
for (i, j), (i2, j2) in zip(s, s[1:]):
assert i <= i2 and j <= j2
assert i2 - i <= max(1, math.ceil(m / p)) and j2 - j <= max(1, math.ceil(n / p))
def test_stable_with_custom_keys() -> None:
class K:
def __init__(self, k: int, tag: str) -> None:
self.k, self.tag = k, tag
def __lt__(self, other: "K") -> bool:
return self.k < other.k
a = [K(1, "a0"), K(2, "a1"), K(2, "a2")]
b = [K(2, "b0"), K(3, "b1")]
out, _ = shiloach_vishkin_merge(a, b, 2)
assert [x.tag for x in out] == ["a0", "a1", "a2", "b0", "b1"]
What the next level needs
- the tests did not run in this build