Kou-Markowsky-Berman algorithm in Python
A 2-approximation for the Steiner tree problem via an MST of the terminal distance graph.
- File
kou_markowsky_berman.py, 120 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/kou-markowsky-berman/. The authors are listed in prov/97ZJ4F.jsonl.
import heapq
def kou_markowsky_berman(n: int, edges: list[tuple[int, int, float]], terminals: list[int]) -> list[int] | None:
"""Return the indices of a Steiner tree for the terminals, or None.
The vertices are 0 .. n - 1 and edges[i] = (u, v, w) is an undirected
edge with w >= 0. Weights may tie (ties are broken by the vertex and the
edge index); parallel edges and self-loops are allowed. The graph need
not be connected, but the terminals must lie in one component, or the
result is None. The tree spans the terminals, all its leaves are
terminals, and its weight is at most 2 (1 - 1/l) times the optimum,
where l is the number of leaves of an optimal tree.
Cost: s runs of Dijkstra's algorithm with a binary heap, where s is the
number of terminals, and two MSTs; O(s (m + n) log n + s^2 log s) time.
With a Fibonacci heap the Dijkstra runs would take O(s (m + n log n)).
"""
if any(w < 0 for _, _, w in edges):
raise ValueError("edge weights must be non-negative")
term = sorted(set(terminals))
if len(term) <= 1:
return []
adj: list[list[tuple[int, float, int]]] = [[] for _ in range(n)]
for i, (u, v, w) in enumerate(edges):
if u != v:
adj[u].append((v, w, i))
adj[v].append((u, w, i))
runs = [_dijkstra(n, adj, s) for s in term] # @line L1
pairs = [] # @line L2
for a in range(len(term)):
for b in range(a + 1, len(term)):
d = runs[a][0][term[b]]
if d is None:
return None
pairs.append((d, a, b))
tree1 = _kruskal(len(term), [(a, b, d) for d, a, b in pairs]) # @line L3
used: set[int] = set() # @line L4
for j in tree1:
_, a, b = pairs[j]
pred = runs[a][1]
x = term[b]
while x != term[a]:
i = pred[x]
used.add(i)
u, v, _ = edges[i]
x = u if v == x else v
sub = sorted(used)
tree2 = [sub[j] for j in _kruskal(n, [edges[i] for i in sub])] # @line L5
return _prune(n, edges, tree2, set(term)) # @line L6
def _dijkstra(n, adj, s):
"""Distances from s (None if unreachable) and the edge into each vertex."""
dist: list[float | None] = [None] * n
pred = [-1] * n
dist[s] = 0
heap = [(0, s)]
done = [False] * n
while heap:
d, u = heapq.heappop(heap)
if done[u]:
continue
done[u] = True
for v, w, i in adj[u]:
nd = d + w
if dist[v] is None or nd < dist[v]:
dist[v] = nd
pred[v] = i
heapq.heappush(heap, (nd, v))
return dist, pred
def _kruskal(n, edges):
"""Indices of a minimum spanning forest, ties broken by the index."""
parent = list(range(n))
def find(x):
while parent[x] != x:
parent[x] = parent[parent[x]]
x = parent[x]
return x
out = []
for i in sorted(range(len(edges)), key=lambda i: (edges[i][2], i)):
a, b = find(edges[i][0]), find(edges[i][1])
if a != b:
parent[a] = b
out.append(i)
return out
def _prune(n, edges, tree, keep):
"""Delete leaves that are not terminals until none is left."""
alive = set(tree)
deg = [0] * n
at: list[list[int]] = [[] for _ in range(n)]
for i in tree:
u, v, _ = edges[i]
deg[u] += 1
deg[v] += 1
at[u].append(i)
at[v].append(i)
stack = [v for v in range(n) if deg[v] == 1 and v not in keep]
while stack:
v = stack.pop()
if deg[v] != 1:
continue
i = next(i for i in at[v] if i in alive)
alive.discard(i)
u, x, _ = edges[i]
other = x if u == v else u
deg[v] -= 1
deg[other] -= 1
if deg[other] == 1 and other not in keep:
stack.append(other)
return sorted(alive)
test_kou_markowsky_berman.py 131 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/kou-markowsky-berman/. The authors are listed in prov/97ZJ4F.jsonl.
import itertools
import random
import pytest
from kou_markowsky_berman import kou_markowsky_berman
EXAMPLE = [(0, 1, 4), (0, 2, 3), (1, 2, 1), (1, 3, 2), (2, 3, 4), (2, 4, 5), (3, 4, 7), (3, 5, 6), (4, 5, 2)]
def find_root(parent, x):
while parent[x] != x:
x = parent[x]
return x
def connects(n, edges, chosen, terms):
parent = list(range(n))
for i in chosen:
parent[find_root(parent, edges[i][0])] = find_root(parent, edges[i][1])
return len({find_root(parent, t) for t in terms}) <= 1
def is_tree(n, edges, chosen):
parent = list(range(n))
for i in chosen:
a, b = find_root(parent, edges[i][0]), find_root(parent, edges[i][1])
if a == b:
return False
parent[a] = b
used = {x for i in chosen for x in edges[i][:2]}
return len({find_root(parent, x) for x in used}) <= 1
def weight(edges, chosen):
return sum(edges[i][2] for i in chosen)
def optimum(n, edges, terms):
"""Least weight of an edge set that connects the terminals, or None."""
best = None
for r in range(len(edges) + 1):
for sub in itertools.combinations(range(len(edges)), r):
if connects(n, edges, sub, terms):
w = weight(edges, sub)
if best is None or w < best:
best = w
return best
def check(n, edges, terms):
got = kou_markowsky_berman(n, edges, terms)
opt = optimum(n, edges, terms)
if opt is None:
assert got is None
return
s = len(set(terms))
assert got is not None and is_tree(n, edges, got)
assert connects(n, edges, got, terms)
for x in {x for i in got for x in edges[i][:2]}:
if sum(x in edges[i][:2] for i in got) == 1:
assert x in terms
w = weight(edges, got)
assert opt <= w
if s >= 2:
assert w * s <= 2 * (s - 1) * opt
def test_example():
assert kou_markowsky_berman(6, EXAMPLE, [0, 1, 4]) == [0, 2, 5]
assert weight(EXAMPLE, [0, 2, 5]) == 10
assert optimum(6, EXAMPLE, [0, 1, 4]) == 9
def test_example_optimal_case():
got = kou_markowsky_berman(6, EXAMPLE, [0, 3, 5])
assert weight(EXAMPLE, got) == 12 == optimum(6, EXAMPLE, [0, 3, 5])
def test_all_vertices_terminal_gives_mst():
got = kou_markowsky_berman(6, EXAMPLE, list(range(6)))
assert weight(EXAMPLE, got) == 13
def test_star_lower_bound():
# Hub 0, terminals 1 .. 4 with spokes 1 and a path of edges of weight 1.9.
edges = [(0, t, 1) for t in range(1, 5)] + [(1, 2, 1.9), (2, 3, 1.9), (3, 4, 1.9)]
got = kou_markowsky_berman(5, edges, [1, 2, 3, 4])
assert got == [4, 5, 6]
assert abs(weight(edges, got) - 5.7) < 1e-9
def test_small_cases():
assert kou_markowsky_berman(3, [(0, 1, 1)], []) == []
assert kou_markowsky_berman(3, [(0, 1, 1)], [2]) == []
assert kou_markowsky_berman(3, [(0, 1, 1)], [1, 1, 0]) == [0]
assert kou_markowsky_berman(3, [(0, 1, 1)], [0, 2]) is None
def test_zero_weights_parallel_edges_and_loops():
edges = [(0, 1, 0), (0, 1, 0), (1, 1, 0), (1, 2, 3), (0, 2, 3)]
got = kou_markowsky_berman(3, edges, [0, 1, 2])
assert weight(edges, got) == 3 and is_tree(3, edges, got)
def test_negative_weight_rejected():
with pytest.raises(ValueError):
kou_markowsky_berman(2, [(0, 1, -1)], [0, 1])
def test_random_against_brute_force():
rng = random.Random(20261105)
for trial in range(600):
n = rng.randint(1, 6)
m = rng.randint(0, 9)
edges = [(rng.randrange(n), rng.randrange(n), rng.randint(0, 5)) for _ in range(m)]
terms = rng.sample(range(n), rng.randint(0, n))
check(n, edges, terms)
def test_random_connected_larger():
rng = random.Random(5)
for _ in range(60):
n = rng.randint(4, 8)
edges = [(v, rng.randrange(v), rng.randint(1, 9)) for v in range(1, n)]
edges += [(rng.randrange(n), rng.randrange(n), rng.randint(1, 9)) for _ in range(rng.randint(0, 4))]
terms = rng.sample(range(n), rng.randint(2, n))
check(n, edges, terms)
What the next level needs
- the tests did not run in this build