Virtual tree in Python
A compressed tree built on a chosen set of key vertices and their pairwise lowest common ancestors, so queries touch only O(k) nodes.
- File
virtual_tree.py, 81 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/virtual-tree/. The authors are listed in prov/Q53SNA.jsonl.
class VirtualTree:
"""Virtual trees of key sets in a fixed rooted tree.
The tree has the vertices 0 .. n-1, the n - 1 given edges and the root
`root`; it is not weighted. The constructor computes the preorder entry
times, the depths and a binary lifting table in Theta(n log n) time,
without recursion. build(keys) returns the virtual tree of the keys.
"""
def __init__(self, n: int, edges: list[tuple[int, int]], root: int = 0) -> None:
adj: list[list[int]] = [[] for _ in range(n)]
for a, b in edges:
adj[a].append(b)
adj[b].append(a)
self.n = n
self.root = root
self.parent = [-1] * n
self.depth = [0] * n
self.tin = [0] * n
seen = [False] * n
seen[root] = True
stack = [root]
t = 0
while stack: # iterative preorder
v = stack.pop()
self.tin[v] = t
t += 1
for u in reversed(adj[v]):
if not seen[u]:
seen[u] = True
self.parent[u] = v
self.depth[u] = self.depth[v] + 1
stack.append(u)
self.log = max(1, (n - 1).bit_length())
up0 = [p if p >= 0 else v for v, p in enumerate(self.parent)]
self.up = [up0]
for _ in range(1, self.log):
prev = self.up[-1]
self.up.append([prev[prev[v]] for v in range(n)])
def lca(self, u: int, v: int) -> int:
"""Lowest common ancestor by binary lifting, O(log n)."""
if self.depth[u] < self.depth[v]:
u, v = v, u
diff = self.depth[u] - self.depth[v]
j = 0
while diff:
if diff & 1:
u = self.up[j][u]
diff >>= 1
j += 1
if u == v:
return u
for j in range(self.log - 1, -1, -1):
if self.up[j][u] != self.up[j][v]:
u, v = self.up[j][u], self.up[j][v]
return self.up[0][u]
def build(self, keys: list[int]) -> tuple[list[int], list[int]]:
"""Return (nodes, par) for the virtual tree of a non-empty key list.
nodes are the vertices of the virtual tree in preorder; par[i] is the
position in nodes of the virtual parent of nodes[i], or -1 for
nodes[0], the LCA of all keys. The edge from nodes[i] to its parent
stands for a path of depth[nodes[i]] - depth[nodes[par[i]]] edges.
"""
tin = self.tin
ks = sorted(set(keys), key=tin.__getitem__) # @line L1
a = list(ks) # @line L2
for i in range(1, len(ks)): # @line L3
a.append(self.lca(ks[i - 1], ks[i]))
nodes = sorted(set(a), key=tin.__getitem__) # @line L4
pos = {v: i for i, v in enumerate(nodes)}
par = [-1] * len(nodes) # @line L5
for i in range(1, len(nodes)): # @line L6
par[i] = pos[self.lca(nodes[i - 1], nodes[i])] # @line L7
return nodes, par # @line L8
test_virtual_tree.py 116 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/virtual-tree/. The authors are listed in prov/Q53SNA.jsonl.
import random
import sys
sys.dont_write_bytecode = True
from virtual_tree import VirtualTree # noqa: E402
EXAMPLE = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (4, 6), (4, 7)]
def naive(n, edges, root, keys):
"""Close the keys under pairwise LCA by climbing, then link each vertex
to its nearest proper ancestor in the closed set."""
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)
def lca(u, v):
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]
return u
ks = set(keys)
closed = {lca(x, y) for x in ks for y in ks}
link = {}
for v in closed:
u = par[v]
while u != -1 and u not in closed:
u = par[u]
link[v] = u
return closed, link
def as_links(nodes, par):
return {v: (nodes[par[i]] if par[i] >= 0 else -1) for i, v in enumerate(nodes)}
def test_example():
vt = VirtualTree(8, EXAMPLE, 0)
nodes, par = vt.build([5, 3, 6])
assert nodes == [0, 1, 3, 6, 5]
assert as_links(nodes, par) == {0: -1, 1: 0, 3: 1, 6: 1, 5: 0}
total = sum(vt.depth[v] - vt.depth[nodes[par[i]]] for i, v in enumerate(nodes) if par[i] >= 0)
assert total == 6
nodes, par = vt.build([6, 7])
assert nodes == [4, 6, 7] and par == [-1, 0, 0]
def test_single_vertex():
vt = VirtualTree(1, [], 0)
assert vt.build([0]) == ([0], [-1])
assert vt.build([0, 0]) == ([0], [-1])
def test_star():
n = 30
vt = VirtualTree(n, [(0, i) for i in range(1, n)], 0)
nodes, par = vt.build([5, 9, 2])
assert nodes == [0, 2, 5, 9] and par == [-1, 0, 0, 0]
assert vt.build([7]) == ([7], [-1])
def test_long_path():
n = 100_000
vt = VirtualTree(n, [(i, i + 1) for i in range(n - 1)], 0)
rng = random.Random(3)
keys = rng.sample(range(n), 2000)
nodes, par = vt.build(keys)
assert nodes == sorted(keys)
assert par == list(range(-1, len(nodes) - 1))
def test_random_against_naive():
rng = random.Random(20240612)
for _ in range(600):
n = rng.randint(1, 40)
perm = list(range(n))
rng.shuffle(perm)
edges = [(perm[i], perm[rng.randrange(i)]) for i in range(1, n)]
root = rng.randrange(n)
vt = VirtualTree(n, edges, root)
keys = [rng.randrange(n) for _ in range(rng.randint(1, 8))]
nodes, par = vt.build(keys)
closed, link = naive(n, edges, root, keys)
assert set(nodes) == closed and len(nodes) == len(closed)
assert len(nodes) <= 2 * len(set(keys)) - 1
assert as_links(nodes, par) == link
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