Tree Problem Patterns — Comprehensive Guide

def dfs(node, parent, adj):
    # PRE-ORDER: process node before children
    print(f"Enter {node}")

    for child in adj[node]:
        if child == parent:
            continue
        dfs(child, node, adj)

    # POST-ORDER: process node after children
    print(f"Leave {node}")

from collections import deque

def bfs(root, adj):
    dist = [-1] * len(adj)
    dist[root] = 0
    queue = deque([root])
    while queue:
        node = queue.popleft()
        for child in adj[node]:
            if dist[child] == -1:
                dist[child] = dist[node] + 1
                queue.append(child)
    return dist

def dfs_iterative(root, adj):
    parent = [-1] * len(adj)
    order = []  # stores nodes in DFS pre-order
    stack = [(root, -1, False)]

    while stack:
        node, par, visited = stack.pop()
        if visited:
            # POST-ORDER work here
            continue
        parent[node] = par
        order.append(node)
        stack.append((node, par, True))  # for post-order
        for child in adj[node]:
            if child != par:
                stack.append((child, node, False))
    return order, parent

class EulerTour:
    def __init__(self, n, adj, root=0):
        self.tin = [0] * n
        self.tout = [0] * n
        self.timer = 0
        self.parent = [-1] * n
        self._dfs(root, -1, adj)

    def _dfs(self, node, par, adj):
        self.parent[node] = par
        self.tin[node] = self.timer
        self.timer += 1
        for child in adj[node]:
            if child != par:
                self._dfs(child, node, adj)
        self.tout[node] = self.timer - 1

    def is_ancestor(self, u, v):
        """O(1) ancestor check."""
        return self.tin[u] <= self.tin[v] <= self.tout[u]

    def subtree_range(self, node):
        """Range in flat array covering node's subtree."""
        return self.tin[node], self.tout[node]

# Subtree sum query
seg.query(tin[node], tout[node])

# Update node value
seg.update(tin[node], new_value)

# Update all nodes in subtree (lazy segment tree)
seg.range_update(tin[node], tout[node], delta)

class HLD:
    def __init__(self, n, adj, root=0):
        self.n = n
        self.adj = adj
        self.parent = [-1] * n
        self.depth = [0] * n
        self.subtree_size = [1] * n
        self.heavy_child = [-1] * n
        self.chain_head = [0] * n
        self.pos = [0] * n
        self.timer = 0

        self._dfs_size(root, -1)
        self.chain_head[root] = root
        self._dfs_hld(root)

    def _dfs_size(self, node, par):
        self.parent[node] = par
        self.depth[node] = 0 if par == -1 else self.depth[par] + 1
        self.subtree_size[node] = 1
        max_size = 0
        for nb in self.adj[node]:
            if nb == par:
                continue
            self._dfs_size(nb, node)
            self.subtree_size[node] += self.subtree_size[nb]
            if self.subtree_size[nb] > max_size:
                max_size = self.subtree_size[nb]
                self.heavy_child[node] = nb

    def _dfs_hld(self, node):
        self.pos[node] = self.timer
        self.timer += 1
        hc = self.heavy_child[node]
        if hc != -1:
            self.chain_head[hc] = self.chain_head[node]
            self._dfs_hld(hc)
        for nb in self.adj[node]:
            if nb == self.parent[node] or nb == hc:
                continue
            self.chain_head[nb] = nb
            self._dfs_hld(nb)

    def query_path(self, u, v, seg):
        """Sum on path u->v using a segment tree."""
        result = 0
        while self.chain_head[u] != self.chain_head[v]:
            if self.depth[self.chain_head[u]] < self.depth[self.chain_head[v]]:
                u, v = v, u
            result += seg.query(self.pos[self.chain_head[u]], self.pos[u])
            u = self.parent[self.chain_head[u]]
        if self.depth[u] > self.depth[v]:
            u, v = v, u
        result += seg.query(self.pos[u], self.pos[v])
        return result

    def lca(self, u, v):
        """LCA as a bonus from HLD."""
        while self.chain_head[u] != self.chain_head[v]:
            if self.depth[self.chain_head[u]] < self.depth[self.chain_head[v]]:
                u, v = v, u
            u = self.parent[self.chain_head[u]]
        return u if self.depth[u] <= self.depth[v] else v

LOG = 20

class LCA:
    def __init__(self, n, adj, root=0):
        self.depth = [0] * n
        self.up = [[-1] * n for _ in range(LOG)]
        self._dfs(root, -1, adj)
        self._build()

    def _dfs(self, node, par, adj):
        self.up[0][node] = par
        for nb in adj[node]:
            if nb != par:
                self.depth[nb] = self.depth[node] + 1
                self._dfs(nb, node, adj)

    def _build(self):
        for k in range(1, LOG):
            for v in range(len(self.depth)):
                mid = self.up[k-1][v]
                self.up[k][v] = -1 if mid == -1 else self.up[k-1][mid]

    def lca(self, u, v):
        if self.depth[u] < self.depth[v]:
            u, v = v, u
        diff = self.depth[u] - self.depth[v]
        for k in range(LOG):
            if diff & (1 << k):
                u = self.up[k][u]
        if u == v:
            return u
        for k in range(LOG - 1, -1, -1):
            if self.up[k][u] != self.up[k][v]:
                u = self.up[k][u]
                v = self.up[k][v]
        return self.up[0][u]

    def dist(self, u, v):
        return self.depth[u] + self.depth[v] - 2 * self.depth[self.lca(u, v)]

def sum_of_distances(n, adj):
    subtree_size = [1] * n
    dp_down = [0] * n   # sum of distances within subtree
    answer = [0] * n

    # Phase 1: DFS down (post-order)
    def dfs_down(node, par):
        for child in adj[node]:
            if child == par:
                continue
            dfs_down(child, node)
            subtree_size[node] += subtree_size[child]
            dp_down[node] += dp_down[child] + subtree_size[child]
            #                                  ^^^^^^^^^^^^^^^^^^
            #                   each node in child's subtree is 1 edge farther

    # Phase 2: DFS up (pre-order)
    def dfs_up(node, par):
        answer[node] = dp_down[node]  # will add dp_up contribution
        if par != -1:
            # dp_up = parent's answer minus our subtree's contribution
            # then add 1 for each node outside our subtree
            outside = answer[par] - (dp_down[node] + subtree_size[node])
            answer[node] += outside + (n - subtree_size[node])

        for child in adj[node]:
            if child != par:
                dfs_up(child, node)

    dfs_down(0, -1)
    answer[0] = dp_down[0]
    dfs_up(0, -1)
    return answer

def farthest_from_each(n, adj):
    # dp_down[node] = [longest_path, second_longest_path] going downward
    dp_down = [[0, 0] for _ in range(n)]
    dp_down_child = [[-1, -1] for _ in range(n)]  # which child gave longest
    dp_up = [0] * n

    def dfs_down(node, par):
        for child in adj[node]:
            if child == par:
                continue
            dfs_down(child, node)
            d = dp_down[child][0] + 1
            if d > dp_down[node][0]:
                dp_down[node][1] = dp_down[node][0]
                dp_down_child[node][1] = dp_down_child[node][0]
                dp_down[node][0] = d
                dp_down_child[node][0] = child
            elif d > dp_down[node][1]:
                dp_down[node][1] = d
                dp_down_child[node][1] = child

    def dfs_up(node, par):
        for child in adj[node]:
            if child == par:
                continue
            # if child gave the longest path, use second longest
            if dp_down_child[node][0] == child:
                best_down = dp_down[node][1]
            else:
                best_down = dp_down[node][0]
            dp_up[child] = 1 + max(dp_up[node], best_down)
            dfs_up(child, node)

    dfs_down(0, -1)
    dfs_up(0, -1)
    return [max(dp_down[i][0], dp_up[i]) for i in range(n)]

from collections import deque

def tree_diameter(n, adj):
    def bfs_farthest(start):
        dist = [-1] * n
        dist[start] = 0
        queue = deque([start])
        farthest = start
        while queue:
            node = queue.popleft()
            for nb in adj[node]:
                if dist[nb] == -1:
                    dist[nb] = dist[node] + 1
                    queue.append(nb)
                    if dist[nb] > dist[farthest]:
                        farthest = nb
        return farthest, dist

    # Step 1: BFS from any node to find one endpoint
    u, _ = bfs_farthest(0)
    # Step 2: BFS from u to find the other endpoint + diameter
    v, dist = bfs_farthest(u)
    return dist[v]  # diameter

def diameter_dp(n, adj):
    diameter = [0]

    def dfs(node, par):
        first = second = 0  # two longest downward paths
        for child in adj[node]:
            if child == par:
                continue
            d = dfs(child, node) + 1
            if d > first:
                second = first
                first = d
            elif d > second:
                second = d
        diameter[0] = max(diameter[0], first + second)
        return first

    dfs(0, -1)
    return diameter[0]

def tree_center(n, adj):
    """Peel leaves layer by layer until 1-2 nodes remain."""
    if n <= 2:
        return list(range(n))

    degree = [len(adj[i]) for i in range(n)]
    # for trees: adjust for undirected adjacency
    leaves = deque(i for i in range(n) if degree[i] <= 1)
    remaining = n

    while remaining > 2:
        new_leaves = deque()
        for leaf in leaves:
            remaining -= 1
            for nb in adj[leaf]:
                degree[nb] -= 1
                if degree[nb] == 1:
                    new_leaves.append(nb)
        leaves = new_leaves

    return list(leaves)

def find_centroid(node, par, adj, subtree_size, tree_size):
    for child in adj[node]:
        if child == par and subtree_size[child] > tree_size // 2:
            # this subtree is too big, centroid is in there
            return find_centroid(child, node, adj, subtree_size, tree_size)
    return node

def centroid_decomposition(n, adj):
    subtree_size = [0] * n
    removed = [False] * n
    centroid_parent = [-1] * n

    def get_size(node, par):
        subtree_size[node] = 1
        for child in adj[node]:
            if child != par and not removed[child]:
                get_size(child, node)
                subtree_size[node] += subtree_size[child]
        return subtree_size[node]

    def get_centroid(node, par, tree_size):
        for child in adj[node]:
            if child != par and not removed[child]:
                if subtree_size[child] > tree_size // 2:
                    return get_centroid(child, node, tree_size)
        return node

    def build(node, par):
        size = get_size(node, -1)
        centroid = get_centroid(node, -1, size)
        removed[centroid] = True
        centroid_parent[centroid] = par

        for child in adj[centroid]:
            if not removed[child]:
                build(child, centroid)

    build(0, -1)
    return centroid_parent

def count_pairs_distance_k(n, adj, k):
    """Count pairs (u,v) where dist(u,v) = k."""
    removed = [False] * n
    subtree_size = [0] * n
    total = 0

    def get_size(node, par):
        subtree_size[node] = 1
        for child in adj[node]:
            if child != par and not removed[child]:
                get_size(child, node)
                subtree_size[node] += subtree_size[child]

    def get_centroid(node, par, tree_size):
        for child in adj[node]:
            if child != par and not removed[child]:
                if subtree_size[child] > tree_size // 2:
                    return get_centroid(child, node, tree_size)
        return node

    def get_depths(node, par, depth, depths):
        depths.append(depth)
        for child in adj[node]:
            if child != par and not removed[child]:
                get_depths(child, node, depth + 1, depths)

    def solve(node):
        nonlocal total
        get_size(node, -1)
        centroid = get_centroid(node, -1, subtree_size[node])
        removed[centroid] = True

        # count paths passing through centroid
        seen = {0: 1}  # distances from centroid already counted
        for child in adj[centroid]:
            if removed[child]:
                continue
            depths = []
            get_depths(child, centroid, 1, depths)
            # count pairs: one from this subtree, one from previous subtrees
            for d in depths:
                need = k - d
                if need >= 0 and need in seen:
                    total += seen[need]
            # add this subtree's depths to seen
            for d in depths:
                seen[d] = seen.get(d, 0) + 1

        for child in adj[centroid]:
            if not removed[child]:
                solve(child)

    solve(0)
    return total

def distinct_colors(n, adj, colors):
    """For each node, count distinct colors in its subtree."""
    answer = [0] * n

    # Each node owns a set. We'll pass sets up the tree.
    node_set = [None] * n

    def dfs(node, par):
        node_set[node] = {colors[node]}

        for child in adj[node]:
            if child == par:
                continue
            dfs(child, node)

            # small-to-large merge
            if len(node_set[child]) > len(node_set[node]):
                node_set[node], node_set[child] = node_set[child], node_set[node]

            # merge smaller into larger
            node_set[node].update(node_set[child])
            node_set[child] = None  # free memory

        answer[node] = len(node_set[node])

    dfs(0, -1)
    return answer

def subtree_sum(node, par, adj, val):
    total = val[node]
    for child in adj[node]:
        if child != par:
            total += subtree_sum(child, node, adj, val)
    return total

def max_independent_set(n, adj, val, root=0):
    dp = [[0, 0] for _ in range(n)]

    def dfs(node, par):
        dp[node][1] = val[node]
        for child in adj[node]:
            if child == par:
                continue
            dfs(child, node)
            dp[node][0] += max(dp[child][0], dp[child][1])
            dp[node][1] += dp[child][0]

    dfs(root, -1)
    return max(dp[root])

def max_matching(n, adj, root=0):
    dp = [[0, 0] for _ in range(n)]

    def dfs(node, par):
        sum_unmatched = 0
        for child in adj[node]:
            if child == par:
                continue
            dfs(child, node)
            sum_unmatched += max(dp[child][0], dp[child][1])

        dp[node][0] = sum_unmatched

        # try matching node with each child
        for child in adj[node]:
            if child == par:
                continue
            gain = dp[child][0] + 1 - max(dp[child][0], dp[child][1])
            dp[node][1] = max(dp[node][1], sum_unmatched + gain)

    dfs(root, -1)
    return max(dp[root])

def paths_through_node(n, adj, root=0):
    """For each node, find the longest path passing through it."""
    best = [0] * n

    def dfs(node, par):
        first = second = 0
        for child in adj[node]:
            if child == par:
                continue
            d = dfs(child, node) + 1
            if d > first:
                first, second = d, first
            elif d > second:
                second = d
        best[node] = first + second  # longest path through this node
        return first  # longest downward path

    dfs(root, -1)
    return best

def build_virtual_tree(query_nodes, tin, lca_func):
    """
    query_nodes: list of node IDs
    tin: euler tour enter times
    lca_func: function that computes LCA of two nodes
    Returns: adjacency list of virtual tree
    """
    # Step 1: sort by tin
    nodes = sorted(query_nodes, key=lambda x: tin[x])

    # Step 2: add LCAs of consecutive pairs
    all_nodes = list(nodes)
    for i in range(len(nodes) - 1):
        l = lca_func(nodes[i], nodes[i + 1])
        all_nodes.append(l)

    # Step 3: deduplicate and sort by tin
    all_nodes = sorted(set(all_nodes), key=lambda x: tin[x])

    # Step 4: add LCAs from stack-based construction
    stack = [all_nodes[0]]
    vtree = {node: [] for node in all_nodes}

    for i in range(1, len(all_nodes)):
        node = all_nodes[i]
        l = lca_func(node, stack[-1])

        if l != stack[-1]:
            while len(stack) > 1 and tin[stack[-2]] >= tin[l]:
                vtree[stack[-2]].append(stack[-1])
                stack.pop()
            if stack[-1] != l:
                vtree[l].append(stack[-1])
                stack[-1] = l

        stack.append(node)

    while len(stack) > 1:
        vtree[stack[-2]].append(stack[-1])
        stack.pop()

    return vtree, stack[0]  # virtual tree adj + root

def tree_to_prufer(n, adj):
    degree = [len(adj[i]) for i in range(n)]
    sequence = []
    leaf = -1
    # find smallest leaf
    for i in range(n):
        if degree[i] == 1:
            leaf = i
            break

    ptr = leaf
    for _ in range(n - 2):
        # find neighbor of current leaf
        for nb in adj[ptr]:
            if degree[nb] > 0:
                neighbor = nb
                break
        sequence.append(neighbor)
        degree[ptr] = 0
        degree[neighbor] -= 1

        if degree[neighbor] == 1 and neighbor < ptr:
            ptr = neighbor  # optimization: this neighbor is now a leaf
        else:
            ptr += 1
            while ptr < n and degree[ptr] != 1:
                ptr += 1

    return sequence

def prufer_to_tree(sequence):
    n = len(sequence) + 2
    degree = [1] * n
    for x in sequence:
        degree[x] += 1

    edges = []
    ptr = 0
    while degree[ptr] != 1:
        ptr += 1

    leaf = ptr
    for x in sequence:
        edges.append((leaf, x))
        degree[leaf] -= 1
        degree[x] -= 1
        if degree[x] == 1 and x < ptr:
            leaf = x
        else:
            ptr += 1
            while ptr < n and degree[ptr] != 1:
                ptr += 1
            leaf = ptr

    edges.append((leaf, n - 1))
    return edges

from collections import defaultdict

def tree_hash(n, adj, root=0):
    """Compute a canonical hash for each subtree."""
    memo = {}  # tuple -> unique id
    counter = [0]
    node_hash = [0] * n

    def get_id(signature):
        if signature not in memo:
            memo[signature] = counter[0]
            counter[0] += 1
        return memo[signature]

    def dfs(node, par):
        child_hashes = []
        for child in adj[node]:
            if child == par:
                continue
            dfs(child, node)
            child_hashes.append(node_hash[child])
        child_hashes.sort()  # canonical ordering
        node_hash[node] = get_id(tuple(child_hashes))

    dfs(root, -1)
    return node_hash

def are_isomorphic(adj1, adj2):
    n1, n2 = len(adj1), len(adj2)
    if n1 != n2:
        return False

    centers1 = tree_center(n1, adj1)
    centers2 = tree_center(n2, adj2)

    for c1 in centers1:
        h1 = tree_hash(n1, adj1, c1)
        for c2 in centers2:
            h2 = tree_hash(n2, adj2, c2)
            if h1[c1] == h2[c2]:
                return True
    return False