LeetCode 834. 树中距离之和

2021-10-09 HGNULB 约 11 分钟

文章信息 7 个分类 5 个标签

分类

题目描述

✅ 834. 树中距离之和

题意分析

给一棵 $n$ 个节点的无向树（$n-1$ 条边、连通、无环），要求返回一个长度为 $n$ 的数组，第 $i$ 项等于节点 $i$ 到其余所有节点的距离之和。距离就是两点间简单路径上的边数。

题面里最重的信号是「对每个节点都要求一个答案」。如果只问某一个节点，一次深度优先遍历就完事了；问所有节点，就意味着要么接受 $n$ 次遍历的平方代价，要么找到相邻两个答案之间的关系。

规模是 $n \le 3 \times 10^4$。$n^2 = 9 \times 10^8$，明确超时，所以平方做法被封死了，必须找递推关系。

树的两条结构性质会在推导里反复用到：其一，树上任意两点的路径唯一，所以「距离」是良定义的；其二，删掉任意一条边 $(u, v)$，树恰好裂成两个连通块，一块含 $u$、一块含 $v$，两块的大小之和等于 $n$。第二条是整个解法的支点。

输入是边的列表而不是父子关系，所以树是「无根」的。要做遍历就得先随便指定一个根（通常取 0），并在遍历时用父节点参数避免走回头路。

边界上要注意两点：$n = 1$ 时唯一的答案是 $0$；$n$ 可以达到 $3 \times 10^4$ 且树可能退化成一条链，递归深度会顶到三万层，栈空间需要留心。

解法：树形 DP + 换根

核心思路

暴力做法很直接：对每个节点跑一次广度优先或深度优先遍历，把所有距离加起来。单次遍历 $O(n)$，总共 $O(n^2)$，$n = 3 \times 10^4$ 时是九亿次操作，超时。

瓶颈在哪？这 $n$ 次遍历之间存在巨量的重复计算——从节点 $u$ 出发和从它的邻居 $v$ 出发，走过的树结构完全一样，只是每个点到起点的距离整体挪了一格。既然如此，能不能只完整算一次，剩下的靠推？

关键观察：设 $u$ 和 $v$ 是一条边的两个端点。以 $u$ 为根时，把根从 $u$ 挪到 $v$，树上每个节点到根的距离要么加一要么减一，而且分界线恰好就是这条边。具体地，去掉边 $(u,v)$ 后，含 $v$ 的那一侧有 $size[v]$ 个节点，它们到新根 $v$ 的距离都比到 $u$ 少 1；含 $u$ 的那一侧有 $n - size[v]$ 个节点，它们到 $v$ 的距离都比到 $u$ 多 1。于是

换根公式就是 $ans[v] = ans[u] - size[v] + (n - size[v])$。

这个公式把「求所有节点的答案」变成了「求一个节点的答案 + 沿边推广」，总代价降到线性。剩下要准备的只有两样东西：某一个节点的完整答案，以及每条边两侧的子树大小。

先把树以 0 为根定向。定义两个数组，含义必须钉死：$size[u]$ 表示以 $u$ 为根的子树里的节点个数（含 $u$ 自己）；$down[u]$ 表示 $u$ 到它子树内部所有节点的距离之和。两者都能在一次自底向上的后序遍历里算出来：$size[u] = 1 + \sum_{v \in child(u)} size[v]$；$down[u] = \sum_{v \in child(u)} \big(down[v] + size[v]\big)$，其中加上 $size[v]$ 是因为子树 $v$ 里的每一个节点，到 $u$ 的距离都比到 $v$ 多走了 $(u,v)$ 这一条边。

注意根节点的特殊性：0 的子树就是整棵树，所以 $down[0]$ 恰好等于 $ans[0]$，第一个完整答案免费到手。

第二遍做自顶向下的先序遍历：进入节点 $u$ 时保证 $ans[u]$ 已经是全局答案，对每个孩子 $v$ 用换根公式算出 $ans[v]$，再递归下去。代码里把 down 和 ans 复用成同一个数组——第一遍结束时它装的是「子树内距离和」，第二遍自顶向下逐个把它改写成「全局距离和」。这个复用之所以安全，靠的是先序的严格顺序：改写 $v$ 时只读 $ans[u]$，而 $u$ 在更外层已经改写完毕；$v$ 自身的旧值在被覆盖前不再被任何人需要。

于是不变量有两条。第一遍结束后，对每个 $u$，$size[u]$ 和 $down[u]$ 都等于其定义值。第二遍进行中，任何时刻「已经被 dfs2 访问过的节点」的 answer 值都等于该节点到全树所有节点的距离之和。

解题步骤

用邻接表建图，每条边正反各加一次。为什么必须双向：输入是无向边，只加单向会让树变成森林，遍历走不通。

第一遍后序遍历 dfs1(u, parent)：先把 $size[u]$ 置为 1，再依次递归每个不等于 parent 的邻居 $v$，回来后累加 $size[u] \mathrel{+}= size[v]$ 和 $answer[u] \mathrel{+}= answer[v] + size[v]$。为什么要传 parent：无向图的邻接表里父节点也是邻居，不挡住会立刻在两点间来回死循环。

为什么累加项是 answer[v] + size[v] 而不只是 answer[v]：answer[v] 统计的是子树 $v$ 内部各点到 $v$ 的距离，换算到以 $u$ 为起点时，这 $size[v]$ 个点每个都要多走边 $(u,v)$，所以整体加上 $size[v]$。

为什么必须是后序（先递归、后累加）：$size[u]$ 和 $answer[u]$ 的取值依赖所有孩子的结果，孩子没算完就累加会取到全零的初值。

第一遍结束后 $answer[0]$ 已经是节点 0 的最终答案，不需要额外处理。为什么：0 是根，它的子树就是整棵树，「子树内距离和」和「全局距离和」在根上重合。

第二遍先序遍历 dfs2(u, parent)：对每个非父邻居 $v$，先执行 $answer[v] = answer[u] - size[v] + (n - size[v])$，再递归进入 $v$。为什么顺序不能颠倒：递归到 $v$ 内部时会读取 $answer[v]$ 当作全局值使用，先递归就会把还是「子树内距离和」的旧值当成全局值传下去，整棵子树全错。

返回 answer 数组。

以 n = 6、edges = [[0,1],[0,2],[2,3],[2,4],[2,5]] 走一遍：以 0 为根，孩子是 1 和 2；2 的孩子是 3、4、5。第一遍自底向上：叶子 1、3、4、5 的 $size$ 都是 1、$answer$ 都是 0。节点 2 有三个叶子孩子，$size[2] = 1 + 1 + 1 + 1 = 4$，$answer[2] = (0+1) + (0+1) + (0+1) = 3$。节点 0 有孩子 1 和 2，$size[0] = 1 + 1 + 4 = 6$，$answer[0] = (answer[1] + size[1]) + (answer[2] + size[2]) = (0 + 1) + (3 + 4) = 8$。

第二遍自顶向下：从 0 出发。孩子 1，$answer[1] = 8 - size[1] + (6 - size[1]) = 8 - 1 + 5 = 12$。孩子 2，$answer[2] = 8 - size[2] + (6 - size[2]) = 8 - 4 + 2 = 6$，注意这一步把 2 身上原本的 3 覆盖掉了，而 3 已经不再被需要。递归进入 2 后，孩子 3，$answer[3] = 6 - 1 + 5 = 10$；孩子 4 和 5 同理都是 $10$。最终返回 [8, 12, 6, 10, 10, 10]。

验算一下节点 1：它到 0 是 1 步，到 2 是 2 步，到 3、4、5 各 3 步，合计 $1 + 2 + 3 \times 3 = 12$，与推导一致。

代码实现

// 第二次 深度优先搜索 通过换根公式计算其它节点的答案： ans[child] = ans[parent] -。
class Solution {
    private List<Integer>[] graph;
    private int[] size;
    private int[] answer;
    private int n;

    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        this.n = n;
        graph = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            graph[i] = new ArrayList<>();
        }
        for (int[] e : edges) {
            graph[e[0]].add(e[1]);
            graph[e[1]].add(e[0]);
        }

        size = new int[n];
        answer = new int[n];

        dfs1(0, -1);
        dfs2(0, -1);

        return answer;
    }

    private void dfs1(int u, int parent) {
        size[u] = 1;
        for (int v : graph[u]) {
            if (v == parent) {
                continue;
            }

            dfs1(v, u);
            size[u] += size[v];
            answer[u] += answer[v] + size[v];
        }
    }

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

            answer[v] = answer[u] - size[v] + (n - size[v]);
            dfs2(v, u);
        }
    }
}

// 第二次 深度优先搜索 通过换根公式计算其它节点的答案： ans[child] = ans[parent] -。
func sumOfDistancesInTree(n int, edges [][]int) []int {
    graph := make([][]int, n)
    for _, e := range edges {
        u := e[0]
        v := e[1]
        graph[u] = append(graph[u], v)
        graph[v] = append(graph[v], u)
    }

    size := make([]int, n)
    answer := make([]int, n)

    var dfs1 func(u int, parent int)
    dfs1 = func(u int, parent int) {
        size[u] = 1
        for _, v := range graph[u] {
            if v == parent {
                continue
            }
            dfs1(v, u)
            size[u] += size[v]
            answer[u] += answer[v] + size[v]
        }
    }

    var dfs2 func(u int, parent int)
    dfs2 = func(u int, parent int) {
        for _, v := range graph[u] {
            if v == parent {
                continue
            }
            answer[v] = answer[u] - size[v] + (n - size[v])
            dfs2(v, u)
        }
    }

    dfs1(0, -1)
    dfs2(0, -1)

    return answer
}

复杂度分析

时间复杂度：$O(n)$，其中 $n$ 是节点数。建图遍历 $n-1$ 条边；两次深度优先遍历各访问每个节点一次、每条边两次（正反各一），加起来仍是 $O(n)$；换根公式本身是常数级运算。

空间复杂度：$O(n)$。邻接表存 $2(n-1)$ 个有向边项，size 和 answer 各占 $n$；递归栈在链状树上会达到 $O(n)$ 深，与前面同阶。

关键点总结

换根 DP 的固定套路是两遍遍历：第一遍后序，自底向上把「子树内」的信息聚合上来；第二遍先序，自顶向下把「子树外」的信息传下去。凡是「对每个节点都要一个全局答案」的树上问题，都值得先往这个模板上套。

换根公式的推导支点是「删掉一条边把树分成两块，大小之和为 $n$」。从 $u$ 挪到 $v$，$size[v]$ 个点各近一步、$n - size[v]$ 个点各远一步，一减一加就是全部变化量。能当场推出来比背下来重要得多。

子树聚合时容易漏掉「跨越这条边的额外一步」。$answer[u] \mathrel{+}= answer[v] + size[v]$ 里的 $size[v]$ 就是这一步的总代价，漏了它答案会小得离谱。

一个数组承担两种语义（第一遍是子树内距离和，第二遍改写成全局距离和）是节省空间的常见技巧，但它的正确性完全依赖第二遍严格的先序顺序。如果不放心，多开一个数组分开存也不丢分。

无向图的邻接表里父节点也是邻居，必须靠 parent 参数屏蔽。用 visited 数组也可以，但在树上 parent 更省事，也更能表明你清楚这是棵树。

面试视角：先说 $O(n^2)$ 的朴素做法并指出重复计算在哪，再推换根公式，最后才谈两遍遍历的实现顺序。$n = 3 \times 10^4$ 且树可能退化成链，递归深度会顶到三万层，主动提一句「必要时改成显式栈的迭代版」是加分项。

易错点总结

错误写法：换根公式写成 answer[v] = answer[u] + size[v] - (n - size[v])，加减号整体写反。用例 n = 6、edges = [[0,1],[0,2],[2,3],[2,4],[2,5]] 中，节点 2 的答案会算成 $8 + 4 - 2 = 10$，正确值是 $6$。

错误写法：第一遍累加时漏掉跨边的一步，写成 answer[u] += answer[v]。同一用例里 $answer[0]$ 会变成 $0 + 3 = 3$ 而不是 $8$，后续所有换根结果连锁出错。

错误写法：dfs2 里先递归子节点再赋值 answer[v]。递归进入 $v$ 时 $answer[v]$ 还是「子树内距离和」，$v$ 的孩子们会基于这个错误的基准继续换根，用例中节点 3、4、5 会算成 $3 - 1 + 5 = 7$。

错误写法：建图时只加 graph[e[0]].add(e[1])。树被当成有向图，从 0 出发只能到达一部分节点，其余节点的 size 和 answer 全为 0。

错误写法：dfs1 里忘记 size[u] = 1 这一行，只累加孩子。所有 size 恒为 0，换根公式退化成 answer[v] = answer[u] + n，结果完全错乱。

错误写法：递归时不传 parent 也不用 visited。无向邻接表里 $u$ 和 $v$ 互为邻居，第一步就会在两点之间无限往返，直接栈溢出。

错误写法：对每个节点单独跑一次广度优先遍历求距离和。逻辑正确但复杂度是 $O(n^2)$，$n = 3 \times 10^4$ 的用例上必然超时。

错误写法：在 dfs2 里读取的是「进入函数时刚被父亲改写前」的 answer[u]，例如把换根写在循环外先缓存了一个旧值。缓存的时机一旦早于父亲的改写，整棵子树都会基于过时基准计算。

错误写法：忽略链状树的递归深度。三万层递归在默认栈大小下可能溢出，遇到这类数据需要改写成显式栈的迭代遍历，或者在支持的环境里调大栈。

相似题目

题目	难度	考察点
310. 最小高度树	中等	同样要考察「以每个节点为根」的某个度量，但只需找出最优根，用拓扑剥叶从外向内收缩即可，不必真的算出每个节点的值
543. 二叉树的直径	简单	只需要一遍自底向上的子树聚合，答案在合并时顺手取最大值，没有第二遍自顶向下的信息回传
396. 旋转函数	中等	同为「从相邻答案递推下一个答案」的换基思想，只是载体从树的换根变成了数组的循环移位，推导时同样要分清哪些项加一、哪些项减一