Tarjan's subtree disassembly in Python
Detects negative cycles in a shortest-path tree by removing subtrees, often faster than Bellman-Ford in practice.
- File
tarjan_subtree_disassembly.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/tarjan-subtree-disassembly/. The authors are listed in prov/8J5TGF.jsonl.
from collections import deque
def tarjan_subtree_disassembly(
adj: list[list[tuple[int, float]]], s: int
) -> tuple[list[float | None] | None, list[int | None] | None, list[int] | None]:
"""Single-source shortest paths with negative-cycle detection by subtree disassembly.
adj[u] lists the arcs (v, w) out of u, for vertices 0 .. n-1, with any real weights.
Returns (dist, parent, None) when no negative cycle is reachable from s: dist[v] is the
distance (None when v is unreachable) and parent is a shortest-path tree. Returns
(None, None, cycle) when it finds a negative cycle: cycle = [c0, c1, ..., ck] with the
arcs c0 -> c1 -> ... -> ck -> c0, of negative total weight. Worst-case time O(n m).
"""
n = len(adj)
dist: list[float | None] = [None] * n # @line L1: None stands for infinity
parent: list[int | None] = [None] * n
children: list[set[int]] = [set() for _ in range(n)]
in_tree = [False] * n
in_queue = [False] * n
dist[s] = 0
in_tree[s] = True
queue: deque[int] = deque([s])
in_queue[s] = True
while queue: # @line L2
u = queue.popleft() # @line L3
if not in_queue[u]:
continue # a stale entry of a vertex that was removed from the queue
in_queue[u] = False
du = dist[u]
assert du is not None and in_tree[u]
for v, w in adj[u]: # @line L4
dv = dist[v]
if dv is not None and du + w >= dv: # @line L5
continue
# collect the proper descendants of v; meeting u means v is an ancestor of u
removed: list[int] = []
found = v == u
if in_tree[v]:
stack = list(children[v])
while stack and not found:
x = stack.pop()
if x == u:
found = True
removed.append(x)
stack.extend(children[x])
if found: # @line L6
cycle = [u] # @line L7: the tree path from v down to u, then the arc (u, v)
while cycle[-1] != v:
p = parent[cycle[-1]]
assert p is not None
cycle.append(p)
cycle.reverse()
return None, None, cycle
for x in removed: # @line L8
in_tree[x] = False
in_queue[x] = False
parent[x] = None
children[x].clear()
children[v].clear()
dist[v] = du + w # @line L9
old = parent[v]
if old is not None:
children[old].discard(v)
parent[v] = u
children[u].add(v)
in_tree[v] = True
if not in_queue[v]: # @line L10
queue.append(v)
in_queue[v] = True
return dist, parent, None # @line L11
test_tarjan_subtree_disassembly.py 121 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/tarjan-subtree-disassembly/. The authors are listed in prov/8J5TGF.jsonl.
import random
import unittest
from tarjan_subtree_disassembly import tarjan_subtree_disassembly
NAMES = "stxyz"
CLRS_22_4 = [
("s", "t", 6), ("s", "y", 7), ("t", "x", 5), ("t", "y", 8), ("t", "z", -4),
("x", "t", -2), ("y", "x", -3), ("y", "z", 9), ("z", "s", 2), ("z", "x", 7),
]
def build(n, edges):
adj = [[] for _ in range(n)]
for u, v, w in edges:
adj[u].append((v, w))
return adj
def named(edges):
return [(NAMES.index(a), NAMES.index(b), w) for a, b, w in edges]
def bellman_ford(n, edges, s):
"""Independent oracle: n - 1 passes, then one more to see a reachable negative cycle."""
dist = [None] * n
dist[s] = 0
for _ in range(n - 1):
for u, v, w in edges:
if dist[u] is not None and (dist[v] is None or dist[u] + w < dist[v]):
dist[v] = dist[u] + w
for u, v, w in edges:
if dist[u] is not None and (dist[v] is None or dist[u] + w < dist[v]):
return None
return dist
def cycle_weight(adj, cycle):
"""Least weight of the closed walk through the listed vertices, using the cheapest parallel arc."""
total = 0
for i, a in enumerate(cycle):
b = cycle[(i + 1) % len(cycle)]
total += min(w for x, w in adj[a] if x == b)
return total
def reachable(adj, s):
seen, stack = {s}, [s]
while stack:
u = stack.pop()
for v, _ in adj[u]:
if v not in seen:
seen.add(v)
stack.append(v)
return seen
class TestSubtreeDisassembly(unittest.TestCase):
def test_clrs_example(self):
dist, parent, cycle = tarjan_subtree_disassembly(build(5, named(CLRS_22_4)), 0)
self.assertIsNone(cycle)
self.assertEqual(dist, [0, 2, 4, 7, -2])
self.assertEqual([NAMES[p] if p is not None else None for p in parent], [None, "x", "y", "s", "t"])
def test_negative_cycle_example(self):
edges = [e if e[:2] != ("x", "t") else ("x", "t", -6) for e in CLRS_22_4]
dist, parent, cycle = tarjan_subtree_disassembly(build(5, named(edges)), 0)
self.assertIsNone(dist)
self.assertEqual([NAMES[c] for c in cycle], ["x", "t"])
def test_self_loop(self):
_, _, cycle = tarjan_subtree_disassembly(build(2, [(0, 1, 1), (1, 1, -1)]), 0)
self.assertEqual(cycle, [1])
_, _, cycle = tarjan_subtree_disassembly(build(1, [(0, 0, -2)]), 0)
self.assertEqual(cycle, [0])
def test_unreachable_negative_cycle_and_vertices(self):
edges = named(CLRS_22_4) + [(5, 6, -1), (6, 5, -1), (5, 0, 1)]
dist, parent, cycle = tarjan_subtree_disassembly(build(8, edges), 0)
self.assertIsNone(cycle)
self.assertEqual(dist, [0, 2, 4, 7, -2, None, None, None])
self.assertEqual(parent[5:], [None, None, None])
def test_single_vertex_and_zero_cycle(self):
self.assertEqual(tarjan_subtree_disassembly([[]], 0), ([0], [None], None))
dist, parent, cycle = tarjan_subtree_disassembly(build(3, [(0, 1, 1), (1, 2, 0), (2, 1, 0)]), 0)
self.assertEqual((dist, parent, cycle), ([0, 1, 1], [None, 0, 1], None))
def test_random_against_bellman_ford(self):
rng = random.Random(1981)
negative = 0
for _ in range(800):
n = rng.randint(1, 8)
m = rng.randint(0, 3 * n)
low = rng.choice([0, -2, -5])
edges = [(rng.randrange(n), rng.randrange(n), rng.randint(low, 9)) for _ in range(m)]
s = rng.randrange(n)
adj = build(n, edges)
expected = bellman_ford(n, edges, s)
dist, parent, cycle = tarjan_subtree_disassembly(adj, s)
if expected is None:
negative += 1
self.assertIsNotNone(cycle)
self.assertEqual(len(set(cycle)), len(cycle))
self.assertLess(cycle_weight(adj, cycle), 0)
self.assertIn(cycle[0], reachable(adj, s))
else:
self.assertIsNone(cycle)
self.assertEqual(dist, expected)
for v in range(n):
if v != s and dist[v] is not None:
u = parent[v]
self.assertTrue(any(x == v and dist[u] + w == dist[v] for x, w in adj[u]))
self.assertGreater(negative, 50)
if __name__ == "__main__":
unittest.main()
What the next level needs
- the tests did not run in this build