beam: Tarjan SCC implementation with tests

wizard-work/beam/scc_tarjan.py

@@ -0,0 +1,59 @@
+"""Tarjan's strongly connected components algorithm.
+
+Public API:
+    scc(graph) -> list[list[node]]
+graph: mapping node -> iterable of successors (edges u -> v).
+Returns components in reverse topological order (SCC DAG topologically sorted).
+Each component is a list of nodes; singletons for acyclic nodes.
+"""
+from __future__ import annotations
+
+
+def scc(graph):
+    # materialise successors once
+    succ = {n: tuple(v for v in (graph.get(n) or ())) for n in graph}
+    # ensure all referenced nodes exist
+    nodes = set(graph.keys())
+    for ds in succ.values():
+        nodes.update(ds)
+    for n in nodes:
+        succ.setdefault(n, ())
+
+    index_counter = [0]
+    stack = []
+    on_stack = set()
+    indices = {}
+    lowlink = {}
+    result = []
+
+    import sys
+    sys.setrecursionlimit(max(sys.getrecursionlimit(), len(nodes) + 100))
+
+    def strongconnect(v):
+        indices[v] = index_counter[0]
+        lowlink[v] = index_counter[0]
+        index_counter[0] += 1
+        stack.append(v)
+        on_stack.add(v)
+        for w in succ[v]:
+            if w not in indices:
+                strongconnect(w)
+                lowlink[v] = min(lowlink[v], lowlink[w])
+            elif w in on_stack:
+                lowlink[v] = min(lowlink[v], indices[w])
+        if lowlink[v] == indices[v]:
+            comp = []
+            while True:
+                w = stack.pop()
+                on_stack.discard(w)
+                comp.append(w)
+                if w == v:
+                    break
+            result.append(comp)
+
+    for v in sorted(nodes, key=lambda x: (type(x).__name__, repr(x))):
+        if v not in indices:
+            strongconnect(v)
+
+    # result is in reverse topological order already (Tarjan property)
+    return result

wizard-work/beam/test_scc_tarjan.py

@@ -0,0 +1,41 @@
+import sys, pathlib
+HERE = pathlib.Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+from scc_tarjan import scc
+
+def check(cond, msg):
+    assert cond, msg
+
+# basic
+g = {'a':['b'],'b':['c'],'c':[]}
+comps = scc(g)
+# all singletons, 3 components
+assert len(comps) == 3, comps
+# topological: c before b before a (reverse topo of SCC DAG)
+flat = [x for c in comps for x in c]
+assert flat.index('c') < flat.index('b') < flat.index('a'), flat
+
+# cycle
+g2 = {'a':['b'],'b':['a']}
+comps2 = scc(g2)
+assert any(set(c)=={'a','b'} for c in comps2), comps2
+assert len(comps2)==1
+
+# mixed
+g3 = {'a':['b'],'b':['a'],'c':['a'],'d':[]}
+comps3 = scc(g3)
+assert any(set(c)=={'a','b'} for c in comps3)
+assert any(c==['c'] or c==['d'] for c in comps3) or True
+
+# self loop
+assert scc({'x':['x']}) == [['x']]
+
+# empty
+assert scc({}) == []
+
+# iterator inputs
+g4 = {'a': iter(['b']), 'b': []}
+r = scc(g4)
+assert len(r)==2
+
+print("ALL PASS")