LeetCode 834. 树中距离之和

2021-10-09 HGNULB 约 6 分钟

文章信息 7 个分类 · 4 个标签

分类

LeetCode CodeTop 困难 未解决 2-5 字节 华为

标签

动态规划 深度优先搜索 树 图

题目描述

✅ 834. 树中距离之和

题意分析

对树中的每个节点，求它到所有其他节点的边数之和。若从每个节点分别遍历一次，会重复计算大量相同路径；相邻节点的距离和只差一次跨边后的增减，可以先算一个节点，再递推出其余答案。

解法：树形 DP + 换根

核心思路

[!blue]

先以 0 为根，记录 parent 和父节点先于孩子的顺序 order。树中没有环，遍历邻接节点时排除父节点就不会重复访问。将 size[u] 初始化为 1，表示当前只统计节点自身；answer[u] 初始为 0，第一遍用于记录 u 到其原根子树中所有节点的距离和。

逆序处理 order 时，孩子都先于父节点完成统计。设 u 的父节点为 p，将 size[u] 加到 size[p]，并将 answer[u] + size[u] 加到 answer[p]：从 p 到 u 子树中的每个节点，都要先多走 p-u 这一条边。汇总结束后，根的子树就是整棵树，所以 answer[0] 已是全局距离和。

再按正序把全局答案从父节点传给孩子。沿 p-u 从 p 移到 u 后，u 原子树中的 size[u] 个节点都近了 1，其余 n - size[u] 个节点都远了 1，因为树中两点之间只有一条路径。因此 answer[u] = answer[p] - size[u] + (n - size[u])，也就是 answer[p] + n - 2 * size[u]。

正序保证使用父节点时，它的 answer 已经是全树答案，随后覆盖孩子之前的局部距离和。size 始终保留以 0 为根时的定义：它表示删去父子边后孩子一侧的节点数，不需要真的修改树结构或重新计算子树。

解题步骤

1. 建立无向邻接表，从 0 出发生成 parent 和 order，将所有 size 设为 1、answer 设为 0。
2. 从 order 末尾处理到下标 1，把当前子树的大小和距离贡献累加给父节点；根没有父节点，必须跳过。
3. 从下标 1 正向处理，使用 answer[parent[u]] + n - 2 * size[u] 得到每个节点的全局答案。
4. 返回 answer。只有一个节点时，两遍循环都跳过，结果自然为 0。

代码实现

class Solution {
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        List<Integer>[] graph = new ArrayList[n];

        for (int i = 0; i < n; i++) {
            graph[i] = new ArrayList<>();
        }

        for (int[] edge : edges) {
            graph[edge[0]].add(edge[1]);
            graph[edge[1]].add(edge[0]);
        }

        int[] parent = new int[n];

        // 零号为初始根，遍历时排除返回父节点的边
        parent[0] = -1;
        int[] order = new int[n];
        int count = 1;

        for (int i = 0; i < count; i++) {
            int u = order[i];

            for (int v : graph[u]) {
                if (v == parent[u]) {
                    continue;
                }

                parent[v] = u;
                // 先记录父节点再记录孩子，后续可按相反方向汇总
                order[count++] = v;
            }
        }

        int[] size = new int[n];

        Arrays.fill(size, 1);
        int[] answer = new int[n];

        // 逆序保证孩子已经完成子树统计，根不再向上汇总
        for (int i = n - 1; i > 0; i--) {
            int u = order[i];
            int p = parent[u];

            size[p] += size[u];
            // 子树每个节点到父节点的距离都比到孩子多一
            answer[p] += answer[u] + size[u];
        }

        for (int i = 1; i < n; i++) {
            int u = order[i];

            // 换根后子树内距离减一，外部距离加一；大小仍按原根定义
            answer[u] = answer[parent[u]] + n - 2 * size[u];
        }

        return answer;
    }
}

func sumOfDistancesInTree(n int, edges [][]int) []int {
    graph := make([][]int, n)
    for _, edge := range edges {
        graph[edge[0]] = append(graph[edge[0]], edge[1])
        graph[edge[1]] = append(graph[edge[1]], edge[0])
    }

    parent := make([]int, n)
    // 零号为初始根，遍历时排除返回父节点的边
    parent[0] = -1
    order := make([]int, 1, n)
    for i := 0; i < len(order); i++ {
        u := order[i]
        for _, v := range graph[u] {
            if v == parent[u] {
                continue
            }
            parent[v] = u
            // 先记录父节点再记录孩子，后续可按相反方向汇总
            order = append(order, v)
        }
    }

    size := make([]int, n)
    for i := range size {
        size[i] = 1
    }
    answer := make([]int, n)
    // 逆序保证孩子已经完成子树统计，根不再向上汇总
    for i := n - 1; i > 0; i-- {
        u := order[i]
        p := parent[u]
        size[p] += size[u]
        // 子树每个节点到父节点的距离都比到孩子多一
        answer[p] += answer[u] + size[u]
    }
    for i := 1; i < n; i++ {
        u := order[i]
        // 换根后子树内距离减一，外部距离加一；大小仍按原根定义
        answer[u] = answer[parent[u]] + n - 2*size[u]
    }
    return answer
}

复杂度分析

• 时间复杂度：$O(n)$，建图与两遍遍历线性。
• 空间复杂度：$O(n)$，邻接表、父数组、遍历顺序、大小和答案数组，不依赖递归深度。

关键点总结

[!green]

• 逆序用于子树汇总，正序用于把父节点的全局答案传向孩子。
• 换根只更新距离和，不真的改动树的连接。
• 子树大小仍以最初的零号根为准。

易错点总结

[!yellow]

• 汇总时漏掉 size[child]，会忽略父节点到孩子这一条边给每个子树节点增加的距离。
• 没有先算出父节点的全局答案就换根，会读取局部距离和。
• 换根时重算或修改子树大小，会破坏递推式的原始划分。

相似题目

题目	难度	关联与区别
2581. 统计可能的树根数目	困难	同样先按一个根计算，再沿树边换根修正答案，原题调整正确猜测数，本题调整全部节点距离和。

转载与许可

作者 HGNULB

链接 https://hgnulb.github.io/blog/2021/10883574

许可 本博客所有文章除特别声明外，均采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处！