beam: Tarjan articulation points and bridges (multigraph-aware)

wizard-work/beam/tarjan_articulation.py

@@ -0,0 +1,101 @@
+"""Tarjan articulation points and bridges in an undirected graph."""
+from __future__ import annotations
+import sys
+from collections import defaultdict, deque
+
+
+def _build_indexed(adj):
+    adj = {k: list(v) for k, v in adj.items()}
+    nodes = set(adj.keys())
+    for u, nbrs in adj.items():
+        for v in nbrs:
+            nodes.add(v)
+    for n in nodes:
+        adj.setdefault(n, [])
+
+    slots = []
+    for u in nodes:
+        for v in adj[u]:
+            if v != u:
+                slots.append((u, v))
+
+    pair_of = [-1] * len(slots)
+    remaining = defaultdict(deque)
+    for i, (u, v) in enumerate(slots):
+        remaining[(v, u)].append(i)
+    for i, (u, v) in enumerate(slots):
+        if pair_of[i] != -1:
+            continue
+        q = remaining[(u, v)]
+        while q and pair_of[q[0]] != -1:
+            q.popleft()
+        if q:
+            r = q.popleft()
+            pair_of[i] = r
+            pair_of[r] = i
+
+    # shared uid per undirected edge
+    uid_of = [None] * len(slots)
+    uid = 0
+    for i in range(len(slots)):
+        if uid_of[i] is not None:
+            continue
+        if pair_of[i] != -1 and pair_of[i] != i:
+            uid_of[i] = uid
+            uid_of[pair_of[i]] = uid
+            uid += 1
+        else:
+            uid_of[i] = uid
+            uid += 1
+
+    idx_adj = {n: [] for n in nodes}
+    for i, (u, v) in enumerate(slots):
+        idx_adj[u].append((v, uid_of[i]))
+
+    return nodes, idx_adj
+
+
+def articulation_points_and_bridges(adj):
+    adj = {k: list(v) for k, v in adj.items()}
+    nodes, idx_adj = _build_indexed(adj)
+
+    sys.setrecursionlimit(max(10000, len(nodes) + 100))
+    disc = {}
+    low = {}
+    timer = [0]
+    ap = set()
+    bridges = set()
+
+    def dfs(u, parent_uid, is_root):
+        disc[u] = low[u] = timer[0]
+        timer[0] += 1
+        children = 0
+        for v, uid in idx_adj[u]:
+            if uid == parent_uid:
+                continue
+            if v not in disc:
+                children += 1
+                dfs(v, uid, False)
+                low[u] = min(low[u], low[v])
+                if (not is_root) and low[v] >= disc[u]:
+                    ap.add(u)
+                if low[v] > disc[u]:
+                    bridges.add((min(u, v), max(u, v)))
+            else:
+                low[u] = min(low[u], disc[v])
+        if is_root and children > 1:
+            ap.add(u)
+
+    for start in sorted(nodes, key=lambda x: str(x)):
+        if start not in disc:
+            dfs(start, -1, True)
+
+    return ap, bridges
+
+
+def articulation_points(adj):
+    return articulation_points_and_bridges(adj)[0]
+
+
+def bridges(adj):
+    return articulation_points_and_bridges(adj)[1]

wizard-work/beam/test_tarjan_articulation.py

@@ -0,0 +1,20 @@
+import sys
+sys.path.insert(0, '/root/.hermes/wizard-outputs/beam')
+from tarjan_articulation import articulation_points, bridges
+
+adj = {'1': ['2'], '2': ['1', '3'], '3': ['2', '4'], '4': ['3']}
+assert sorted(articulation_points(adj)) == ['2', '3'], articulation_points(adj)
+assert sorted(bridges(adj)) == [('1', '2'), ('2', '3'), ('3', '4')], bridges(adj)
+
+adj = {'a': ['b', 'c'], 'b': ['a', 'c'], 'c': ['a', 'b']}
+assert articulation_points(adj) == set(), articulation_points(adj)
+assert bridges(adj) == set(), bridges(adj)
+
+adj = {'1': ['2', '2'], '2': ['1', '1']}
+assert bridges(adj) == set(), bridges(adj)
+
+adj = {'c': ['a', 'b', 'd'], 'a': ['c'], 'b': ['c'], 'd': ['c']}
+assert articulation_points(adj) == {'c'}, articulation_points(adj)
+
+print("ALL TESTS PASS")
+print("written")