cswork / libraryOpen practice workspace

Algorithm tutorials

Tree Diameter and Center

Find the longest path in a tree with two BFS passes or one bottom-up DP (needed for negative weights and binary-tree variants), and find the center(s) by peeling leaves or taking the middle of a diameter.

Intuition

The diameter of a tree is the longest simple path between any two vertices (counted in edges, or in total weight). A tree can have several diameters, but they all have the same length.

Two facts carry this whole topic:

  • The farthest vertex from anywhere is a diameter endpoint (for non-negative weights). Start at any vertex, walk to the farthest vertex a; from a, the farthest vertex b gives a diameter a-b. Two traversals, done.
  • Every path has a highest vertex. If you root the tree, the longest path "bends" at some vertex u: it goes down into two different child subtrees. So the diameter is the maximum over u of the two longest downward branches from u added together. One bottom-up pass computes it, and unlike the two-pass trick this also works with negative weights.

The center is the vertex whose farthest vertex is as close as possible (minimum eccentricity). Rooting the tree there minimizes its height, which is the "minimum height trees" question. Useful properties:

  • there are one or two centers, and two centers are adjacent;
  • the centers are the middle vertex (or middle two vertices) of any diameter;
  • the radius (the center's eccentricity) is ceil(D / 2) for a diameter of D edges.

When to use it in interviews / OAs

  • "Longest path between any two nodes", "diameter of a binary tree", "maximum path sum" (the bend-at-the-top DP with node values).
  • "Root the tree to minimize height" / "minimum height trees": centers.
  • "Farthest node from each node", "time for a signal to reach everybody if it starts at the best node".
  • "Connect two trees with one edge so the result has the smallest diameter": connect the centers.
  • "Minimum number of moves to visit all nodes and/or return": often 2 * (edges) - diameter.

Template

Two BFS passes (unweighted, or non-negative weights with a weighted traversal)

import sys
from collections import deque
from dataclasses import dataclass


def farthest(adj: list[list[int]], src: int) -> tuple[int, list[int], list[int]]:
    dist = [-1] * len(adj)
    parent = [-1] * len(adj)
    dist[src] = 0
    q = deque([src])
    last = src
    while q:
        u = q.popleft()
        last = u  # BFS pops in order of distance, so the last one is farthest
        for v in adj[u]:
            if dist[v] == -1:
                dist[v] = dist[u] + 1
                parent[v] = u
                q.append(v)
    return last, dist, parent


def diameter_path(n: int, edges: list[tuple[int, int]]) -> list[int]:
    adj: list[list[int]] = [[] for _ in range(n)]
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)
    a, _, _ = farthest(adj, 0)
    b, _, parent = farthest(adj, a)
    path = [b]
    while path[-1] != a:
        path.append(parent[path[-1]])
    return path  # len(path) - 1 edges


def centers_from_path(path: list[int]) -> list[int]:
    L = len(path) - 1  # diameter in edges
    return [path[L // 2]] if L % 2 == 0 else [path[L // 2], path[L // 2 + 1]]

One bottom-up pass (any weights)

down[u] is the heaviest path that starts at u and goes down into its subtree (0 means "just u"). The best path that bends at u joins the two best branches.

def diameter_dp(n: int, edges: list[tuple[int, int, int]]) -> int:
    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))
    parent, order, stack = [-1] * n, [], [0]
    while stack:  # iterative preorder from root 0
        u = stack.pop()
        order.append(u)
        for v, _ in adj[u]:
            if v != parent[u]:
                parent[v] = u
                stack.append(v)
    down = [0] * n
    best = 0
    for u in reversed(order):  # children are finished before their parent
        for v, w in adj[u]:
            if v != parent[u]:
                branch = down[v] + w
                best = max(best, down[u] + branch)  # bend at u: best branch so far + this one
                down[u] = max(down[u], branch)
    return best

The binary-tree version counts edges and uses each node's height:

@dataclass
class TreeNode:
    val: int
    left: "TreeNode | None" = None
    right: "TreeNode | None" = None


def diameter_of_binary_tree(root: TreeNode | None) -> int:
    sys.setrecursionlimit(10**5)
    best = 0

    def height(node: TreeNode | None) -> int:  # nodes on the longest downward path
        nonlocal best
        if not node:
            return 0
        lh, rh = height(node.left), height(node.right)
        best = max(best, lh + rh)  # edges of the path bending at node
        return 1 + max(lh, rh)

    height(root)
    return best

Centers by peeling leaves (minimum height trees)

Remove all current leaves layer by layer; the last one or two vertices standing are the centers.

def tree_centers(n: int, edges: list[tuple[int, int]]) -> list[int]:
    if n == 1:
        return [0]
    adj: list[list[int]] = [[] for _ in range(n)]
    deg = [0] * n
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)
        deg[u] += 1
        deg[v] += 1
    leaves = [v for v in range(n) if deg[v] == 1]
    remaining = n
    while remaining > 2:
        remaining -= len(leaves)
        nxt = []
        for u in leaves:
            for v in adj[u]:
                deg[v] -= 1
                if deg[v] == 1:  # v just became a leaf
                    nxt.append(v)
        leaves = nxt
    return sorted(leaves)


def merged_diameter(d1: int, d2: int) -> int:
    # join two trees (diameters d1, d2 in edges) by one edge between their centers
    return max(d1, d2, (d1 + 1) // 2 + (d2 + 1) // 2 + 1)

Complexity

Every routine above is O(n) time and O(n) memory: two BFS passes, one bottom-up pass, or one peeling pass where each vertex leaves once.

Common pitfalls

  • Two passes with negative weights. "The farthest vertex is a diameter endpoint" needs non-negative weights. Use the DP.
  • Edges vs. nodes. "Diameter of a binary tree" usually counts edges; "longest path" problems sometimes count nodes. Off by one either way.
  • Initializing down with negative infinity in the DP forbids the single-vertex path and breaks trees with only negative edges; 0 means "stop here".
  • n = 1 and n = 2 in leaf peeling: one vertex has degree 0 (it is the center), and with two vertices both are centers.
  • Recursion depth on path-shaped trees. The DP above is iterative; the binary-tree version needs a raised limit or an explicit stack for very deep trees.
  • Weighted center. With edge weights the eccentricity minimizer can lie in the middle of an edge; vertex-only answers then need the general DP (compute down and up distances for every vertex).

Interview variations

  • Binary tree maximum path sum. Same bend-at-the-top DP with node values; clamp negative branches to 0.
  • Longest path with different adjacent characters / longest univalue path. The DP, but a branch only extends when the labels allow it.
  • Minimum height trees. Leaf peeling, or the middle of a diameter.
  • Merge two trees with minimum diameter. max(d1, d2, ceil(d1/2) + ceil(d2/2) + 1).
  • Visit all nodes. Starting anywhere and not returning, a tree walk covering every edge needs 2(n - 1) - D steps (D = diameter in edges).
  • Farthest node for every node. Two BFS from the diameter endpoints a and b; the answer for v is max(dist_a[v], dist_b[v]).