diff --git a/wizard-work/beam/tarjan_articulation.py b/wizard-work/beam/tarjan_articulation.py new file mode 100644 index 0000000..7a81769 --- /dev/null +++ b/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] \ No newline at end of file diff --git a/wizard-work/beam/test_tarjan_articulation.py b/wizard-work/beam/test_tarjan_articulation.py new file mode 100644 index 0000000..ba656aa --- /dev/null +++ b/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") \ No newline at end of file