Tree diameter by double sweep in Python
Search from any node to the farthest node, then from there; the second distance is the diameter.
- File
tree_diameter_double_sweep.py, 50 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/tree-diameter-double-sweep/. The authors are listed in prov/EZ41AA.jsonl.
Edge = tuple[int, int, int]
def distances(adj: list[list[tuple[int, int]]], s: int) -> tuple[list[int], list[int]]:
"""Distances and parents from s in a tree, by an iterative depth-first search.
adj[u] lists the pairs (v, w) of the neighbours of u and the edge lengths.
The parent of s is -1. Time Theta(n), no recursion.
"""
n = len(adj)
dist = [0] * n
par = [-1] * n
seen = [False] * n
seen[s] = True # @line L9
stack = [s]
while stack: # @line L10
u = stack.pop() # @line L11
for v, w in adj[u]: # @line L12
if not seen[v]:
seen[v] = True
dist[v] = dist[u] + w # @line L13
par[v] = u
stack.append(v) # @line L14
return dist, par # @line L15
def tree_diameter(n: int, edges: list[Edge]) -> tuple[int, list[int]]:
"""Return (D, path): the diameter of a tree and one diametral path.
The tree is free, with the vertices 0 .. n-1 (n >= 1) and the n - 1 edges
(u, v, w), where w >= 0 is the length. For an unweighted tree use w = 1.
The path lists the vertices from one end a to the other end b. Time and
extra space Theta(n). The result can be wrong if some w is negative.
"""
adj: list[list[tuple[int, int]]] = [[] for _ in range(n)]
for u, v, w in edges:
adj[u].append((v, w))
adj[v].append((u, w))
dist, _ = distances(adj, 0) # @line L1
a = max(range(n), key=dist.__getitem__) # @line L2
dist, par = distances(adj, a) # @line L3
b = max(range(n), key=dist.__getitem__) # @line L4
path = [b] # @line L5
while path[-1] != a: # @line L6
path.append(par[path[-1]]) # @line L7
path.reverse()
return dist[b], path # @line L8
test_tree_diameter_double_sweep.py 81 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/tree-diameter-double-sweep/. The authors are listed in prov/EZ41AA.jsonl.
import random
from tree_diameter_double_sweep import tree_diameter
def brute(n: int, edges: list[tuple[int, int, int]]) -> int:
# Distances by repeated relaxation over the edge list from every source.
best = 0
for s in range(n):
d: list[int | None] = [None] * n
d[s] = 0
changed = True
while changed:
changed = False
for u, v, w in edges:
for x, y in ((u, v), (v, u)):
if d[x] is not None and d[y] is None:
d[y] = d[x] + w
changed = True
best = max(best, max(x for x in d if x is not None))
return best
def path_length(edges: list[tuple[int, int, int]], path: list[int]) -> int:
w = {}
for u, v, x in edges:
w[(u, v)] = w[(v, u)] = x
assert len(set(path)) == len(path)
return sum(w[(path[i], path[i + 1])] for i in range(len(path) - 1))
def test_example() -> None:
edges = [(0, 1, 1), (0, 2, 1), (1, 3, 1), (1, 4, 1), (2, 5, 1), (4, 6, 1), (4, 7, 1)]
d, path = tree_diameter(8, edges)
assert d == 5
assert path in ([6, 4, 1, 0, 2, 5], [5, 2, 0, 1, 4, 6], [7, 4, 1, 0, 2, 5], [5, 2, 0, 1, 4, 7])
def test_weighted_example() -> None:
edges = [(0, 1, 2), (1, 2, 7), (1, 3, 3), (0, 4, 4), (4, 5, 1), (4, 6, 5)]
d, path = tree_diameter(7, edges)
assert d == 18
assert path_length(edges, path) == 18
def test_single_vertex() -> None:
assert tree_diameter(1, []) == (0, [0])
def test_path_long() -> None:
n = 100_000
edges = [(i, i + 1, 1) for i in range(n - 1)]
d, path = tree_diameter(n, edges)
assert d == n - 1
assert path == list(range(n - 1, -1, -1))
def test_star() -> None:
n = 30
d, path = tree_diameter(n, [(0, i, 1) for i in range(1, n)])
assert d == 2 and path[1] == 0 and len(path) == 3
assert tree_diameter(2, [(0, 1, 5)])[0] == 5
def test_zero_weights() -> None:
assert tree_diameter(3, [(0, 1, 0), (1, 2, 0)])[0] == 0
def test_random_against_brute_force() -> None:
rng = random.Random(9124)
for _ in range(600):
n = rng.randint(1, 25)
perm = list(range(n))
rng.shuffle(perm)
edges = [(perm[i], perm[rng.randrange(i)], rng.randint(0, 9)) for i in range(1, n)]
d, path = tree_diameter(n, edges)
assert d == brute(n, edges)
assert path_length(edges, path) == d
What the next level needs
- the tests did not run in this build