【算法学习笔记】换根DP -- java代码实现

1. 问题描述

给定一个无向、连通的树。树中有 n 个标记为 0…n-1 的节点以及 n-1 条边 。
给定整数 n 和数组 edges ， edges[i] = [$a_i$, $b_i$]表示树中的节点 $a_i$ 和 $b_i$ 之间有一条边。
返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。

1.1 例子说明

输入: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释: 树如图所示。我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。

2. 思路

一开始看到往往想所有节点都做一遍根，然后各自深度优先遍历一下，这样做的时间复杂度$O(n^2)$，那能不能有什么方法优化一下呢，这时候就引出要讲的换根DP算法：

上面的图片分别是以0和2作为根节点，然后到其余各个节点的距离和，我们可以发现当改变根时候只有新根及其子树会减少路径，而其余节点会相应增加路径（下移了）
所以我们可以得到 ans[2] = ans[0] + 2 - 4 ，此时我们就可以省去很多的计算时间，但是需要注意的是这种方法需要得到每一个节点的子树！！，这个我们可以在第一遍dfs的时候做

---

将上式进行一个抽象我们就可以得到递推方程

1. 假设当前根节点是x，而下一个做根节点的是y。
2. 假设第 i 个节点的子树记录在size数组中（非i节点的所有子树就是 n-size[i] ）。则有：

ans[y] = ans[x] +n - size[y] - size[y]
           ~~~~~~~~~~            = ans[x] +n - 2*size[y] 即为核心公式！！

2.2 算法流程

1. 从 0 出发进行 DFS，累加 0 到每一个点的距离，得到初始$ans[0]$
2. DFS同时计算子树的大小 $size[i]$
3. 第二遍从 0 出发 DFS，设 $y$ 是 $x$ 的儿子，那么有：$ans[y]=ans[x]+n-2\ast size[y]$
4. 递归得到每一个$ans[i]$

3. 代码实现

class Solution {

    private List<Integer>[] g;

    private int[] ans, size;



    public int[] sumOfDistancesInTree(int n, int[][] edges) {

        g = new ArrayList[n]; // g[x] 表示 x 的所有邻居

        Arrays.setAll(g, e -> new ArrayList<>());

        for (var e : edges) {

            int x = e[0], y = e[1];

            g[x].add(y);

            g[y].add(x);

        }

        ans = new int[n];

        size = new int[n];

        dfs(0, -1, 0); // 0 没有父节点

        reroot(0, -1); // 0 没有父节点

        return ans;

    }



    private void dfs(int x, int fa, int depth) {

        ans[0] += depth; // depth 为 0 到 x 的距离

        size[x] = 1;

        for (int y : g[x]) { // 遍历 x 的邻居 y

            if (y != fa) { // 避免访问父节点

                dfs(y, x, depth + 1); // x 是 y 的父节点

                size[x] += size[y]; // 累加 x 的儿子 y 的子树大小

            }

        }

    }



    private void reroot(int x, int fa) {

        for (int y : g[x]) { // 遍历 x 的邻居 y

            if (y != fa) { // 避免访问父节点

                ans[y] = ans[x] + g.length - 2 * size[y];

                reroot(y, x); // x 是 y 的父节点

            }

        }

    }

}