Tarjan's offline LCA algorithm in Python
Answer a batch of LCA queries offline in near-linear time using DFS and union-find.
- File
tarjan_offline_lca.py, 70 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-offline-lca/. The authors are listed in prov/71EN86.jsonl.
def offline_lca(
n: int,
edges: list[tuple[int, int]],
queries: list[tuple[int, int]],
root: int = 0,
) -> list[int]:
"""Answer the LCA queries of a rooted tree offline, with Tarjan's method.
The tree has the vertices 0 .. n-1, the n - 1 given edges, and the root
`root`. It is not weighted. queries[i] = (u, v) asks for the lowest common
ancestor of u and v; the result holds the answers in the same order.
Time O((n + q) alpha(n)) with q = len(queries); no recursion.
"""
adj: list[list[int]] = [[] for _ in range(n)]
for a, b in edges:
adj[a].append(b)
adj[b].append(a)
uf = list(range(n)) # @line L1
rank = [0] * n
anc = list(range(n))
done = [False] * n
def find(x: int) -> int:
r = x
while uf[r] != r:
r = uf[r]
while uf[x] != r: # path compression
uf[x], x = r, uf[x]
return r
def union(x: int, y: int) -> None:
x, y = find(x), find(y)
if x == y:
return
if rank[x] < rank[y]:
x, y = y, x
uf[y] = x
if rank[x] == rank[y]:
rank[x] += 1
qs: list[list[tuple[int, int]]] = [[] for _ in range(n)]
for i, (u, v) in enumerate(queries): # @line L2
qs[u].append((v, i))
qs[v].append((u, i))
ans = [-1] * len(queries)
nxt = [0] * n
stack = [(root, -1)] # @line L3
while stack: # @line L4
v, p = stack[-1] # @line L5
a = adj[v]
while nxt[v] < len(a) and a[nxt[v]] == p:
nxt[v] += 1
if nxt[v] < len(a): # @line L6
c = a[nxt[v]]
nxt[v] += 1
stack.append((c, v))
else: # @line L7
stack.pop() # @line L8
done[v] = True
for w, i in qs[v]: # @line L9
if done[w]:
ans[i] = anc[find(w)]
if p != -1: # @line L10
union(p, v)
anc[find(p)] = p
return ans # @line L11
test_tarjan_offline_lca.py 94 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/tarjan-offline-lca/. The authors are listed in prov/71EN86.jsonl.
import random
import sys
sys.dont_write_bytecode = True
from tarjan_offline_lca import offline_lca # noqa: E402
EXAMPLE = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (4, 6), (4, 7)]
def naive_lca(n, edges, queries, root):
adj = [[] for _ in range(n)]
for a, b in edges:
adj[a].append(b)
adj[b].append(a)
par = [-1] * n
depth = [0] * n
seen = [False] * n
seen[root] = True
order = [root]
for v in order:
for u in adj[v]:
if not seen[u]:
seen[u] = True
par[u] = v
depth[u] = depth[v] + 1
order.append(u)
out = []
for u, v in queries:
while depth[u] > depth[v]:
u = par[u]
while depth[v] > depth[u]:
v = par[v]
while u != v:
u, v = par[u], par[v]
out.append(u)
return out
def random_tree(n, rng):
perm = list(range(n))
rng.shuffle(perm)
return [(perm[i], perm[rng.randrange(i)]) for i in range(1, n)]
def test_example():
q = [(6, 3), (7, 5), (6, 7), (4, 6), (2, 2), (3, 0)]
assert offline_lca(8, EXAMPLE, q, 0) == [1, 0, 4, 4, 2, 0]
assert offline_lca(8, EXAMPLE, [(6, 3), (7, 5)], 4) == [4, 4]
def test_single_vertex():
assert offline_lca(1, [], [(0, 0)], 0) == [0]
assert offline_lca(1, [], [], 0) == []
def test_star():
n = 50
edges = [(0, i) for i in range(1, n)]
q = [(i, j) for i in range(n) for j in range(n)]
assert offline_lca(n, edges, q, 0) == [i if i == j else 0 for i, j in q]
# rooted at the leaf 1: the centre 0 is the LCA of two other leaves
assert offline_lca(n, edges, [(2, 3), (1, 5), (0, 7)], 1) == [0, 1, 0]
def test_long_path():
n = 100_000
edges = [(i, i + 1) for i in range(n - 1)]
rng = random.Random(7)
q = [(rng.randrange(n), rng.randrange(n)) for _ in range(1000)]
assert offline_lca(n, edges, q, 0) == [min(u, v) for u, v in q]
assert offline_lca(n, edges, q, n - 1) == [max(u, v) for u, v in q]
def test_random_against_naive():
rng = random.Random(20240611)
for _ in range(600):
n = rng.randint(1, 40)
edges = random_tree(n, rng)
root = rng.randrange(n)
q = [(rng.randrange(n), rng.randrange(n)) for _ in range(rng.randint(0, 30))]
assert offline_lca(n, edges, q, root) == naive_lca(n, edges, q, root)
if __name__ == "__main__":
test_example()
test_single_vertex()
test_star()
test_long_path()
test_random_against_naive()
print("ok")
What the next level needs
- the tests did not run in this build