LCA by binary lifting in Python
Precompute 2^k ancestors, then answer each query in O(log n).
- File
lca_binary_lifting.py, 76 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/lca-binary-lifting/. The authors are listed in prov/T1VPR2.jsonl.
from collections import deque
class LcaBinaryLifting:
"""Lowest common ancestors in a static rooted tree by binary lifting.
The tree has the vertices 0 .. n-1 (n >= 1), is unweighted, and is given
by its n - 1 undirected edges and a root. Preprocessing takes
O(n log n) time and memory; each query takes O(log n) time.
"""
def __init__(self, n: int, edges: list[tuple[int, int]], root: int = 0) -> None:
if n < 1 or len(edges) != n - 1 or not 0 <= root < n:
raise ValueError("a tree needs n >= 1 vertices, n - 1 edges and a root in range")
adj: list[list[int]] = [[] for _ in range(n)]
for a, b in edges:
adj[a].append(b)
adj[b].append(a)
parent = [-1] * n
depth = [0] * n
parent[root] = root
seen = [False] * n
seen[root] = True
queue = deque([root]) # @line L1
while queue:
x = queue.popleft()
for y in adj[x]:
if not seen[y]:
seen[y] = True
parent[y] = x
depth[y] = depth[x] + 1
queue.append(y)
if not all(seen):
raise ValueError("the edges do not connect all vertices")
self.n = n
self.depth = depth
self.K = max(1, (n - 1).bit_length()) # @line L2
up = [parent[:]] # @line L3
for k in range(1, self.K): # @line L4
prev = up[k - 1]
up.append([prev[prev[v]] for v in range(n)])
self.up = up
def kth_ancestor(self, v: int, d: int) -> int:
"""The ancestor of v at distance d, for 0 <= d <= depth(v)."""
if not 0 <= d <= self.depth[v]:
raise ValueError("d must be between 0 and the depth of v")
k = 0
while d:
if d & 1:
v = self.up[k][v]
d >>= 1
k += 1
return v
def lca(self, u: int, v: int) -> int:
"""The deepest vertex that is an ancestor of both u and v."""
up = self.up
if self.depth[u] < self.depth[v]: # @line L5
u, v = v, u
d = self.depth[u] - self.depth[v] # @line L6
u = self.kth_ancestor(u, d) # @line L7
if u == v: # @line L8
return u
for k in range(self.K - 1, -1, -1): # @line L9
if up[k][u] != up[k][v]: # @line L10
u = up[k][u]
v = up[k][v]
return up[0][u] # @line L11
def distance(self, u: int, v: int) -> int:
"""The number of edges on the path from u to v."""
return self.depth[u] + self.depth[v] - 2 * self.depth[self.lca(u, v)]
test_lca_binary_lifting.py 94 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/lca-binary-lifting/. The authors are listed in prov/T1VPR2.jsonl.
import random
import pytest
from lca_binary_lifting import LcaBinaryLifting
PARENT = [-1, 0, 1, 2, 3, 4, 2, 6, 0, 8, 3, 7]
def tree_edges(parent: list[int]) -> list[tuple[int, int]]:
return [(v, p) for v, p in enumerate(parent) if p >= 0]
def naive_lca(parent: list[int], u: int, v: int) -> int:
seen = set()
while u != -1:
seen.add(u)
u = parent[u]
while v not in seen:
v = parent[v]
return v
def test_example() -> None:
t = LcaBinaryLifting(12, tree_edges(PARENT), 0)
assert t.K == 4
assert t.up[1] == [0, 0, 0, 1, 2, 3, 1, 2, 0, 0, 2, 6]
assert t.up[2] == [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1]
assert t.lca(5, 11) == 2
assert t.lca(10, 9) == 0
assert t.lca(4, 10) == 3
assert t.lca(7, 2) == 2
assert t.lca(6, 6) == 6
assert t.kth_ancestor(11, 5) == 0
assert t.kth_ancestor(5, 3) == 2
assert t.distance(5, 11) == 6
def test_single_vertex() -> None:
t = LcaBinaryLifting(1, [], 0)
assert t.lca(0, 0) == 0
assert t.kth_ancestor(0, 0) == 0
with pytest.raises(ValueError):
t.kth_ancestor(0, 1)
def test_star() -> None:
n = 50
t = LcaBinaryLifting(n, [(0, i) for i in range(1, n)], 0)
assert all(t.lca(i, j) == (i if i == j else 0) for i in range(1, n) for j in range(1, n))
t2 = LcaBinaryLifting(n, [(0, i) for i in range(1, n)], 7)
assert t2.lca(3, 9) == 0
assert t2.lca(0, 9) == 0
assert t2.lca(7, 9) == 7
def test_long_path() -> None:
n = 100_000
t = LcaBinaryLifting(n, [(i, i + 1) for i in range(n - 1)], 0)
assert t.lca(n - 1, 12345) == 12345
assert t.kth_ancestor(n - 1, n - 1) == 0
assert t.distance(0, n - 1) == n - 1
# n = 2^j + 1 needs a jump of 2^j
m = 2**10 + 1
s = LcaBinaryLifting(m, [(i, i + 1) for i in range(m - 1)], 0)
assert s.kth_ancestor(m - 1, m - 1) == 0
assert s.lca(m - 1, 1) == 1
def test_bad_input() -> None:
with pytest.raises(ValueError):
LcaBinaryLifting(0, [], 0)
with pytest.raises(ValueError):
LcaBinaryLifting(4, [(0, 1), (2, 3), (0, 1)], 0)
def test_against_brute_force() -> None:
rng = random.Random(7)
for _ in range(600):
n = rng.randint(1, 40)
parent = [-1] + [rng.randrange(i) for i in range(1, n)]
perm = list(range(n))
rng.shuffle(perm)
relabeled = [-1] * n
for v in range(1, n):
relabeled[perm[v]] = perm[parent[v]]
root = perm[0]
t = LcaBinaryLifting(n, tree_edges(relabeled), root)
for _ in range(20):
u, v = rng.randrange(n), rng.randrange(n)
assert t.lca(u, v) == naive_lca(relabeled, u, v)
What the next level needs
- the tests did not run in this build