树形dp换根

模板题：834. 树中距离之和

给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。

第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。

返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

示例 1:

输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释: 
如下为给定的树的示意图：
  0
 / \
1   2
   /|\
  3 4 5

我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) 
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。

状态转移：

 dp[u] 表示以 u 为根的子树，它的所有子节点到它的距离之和，同时定义 sz[u] 表示以 u 为根的子树的节点数量,son[u] 表示 u 的所有后代节点集合。

对每个节点都做一次如上的树形动态规划，这样时间复杂度即为 O(N^2),通过换根得方法可以对复杂度进行优化

假设 u 的某个后代节点为 v，如果要算 v 的答案，本来我们要以 v 为根来进行一次树形动态规划。但是利用已有的信息，我们可以考虑树的形态做一次改变，让 v 换到根的位置，u 变为其孩子节点，同时维护原有的 dp 信息。在这一次的转变中，我们观察到除了 u 和 v 的 dp 值，其他节点的 dp 值都不会改变，因此只要更新 dp[u] 和 dp[v] 的值即可。

那么我们来看 v 换到根的位置的时候怎么利用已有信息求出 dp[u] 和 dp[v] 的值。重新回顾第一次树形动态规划的转移方程，我们可以知道当 u 变为 v 的孩子的时候 v 不在 u 的后代集合 son[u] 中了，因此此时 dp[u] 需要减去 v 的贡献，即

同时 sz[u] 也要相应减去 sz[v]。

而 v 的后代节点集合中多出了 u，因此dp[v] 的值要由 u 更新上来，即

同时 sz[v] 也要相应加上 sz[u]。

至此我们完成了一次「换根」操作，在 O(1) 的时间内维护了dp 的信息，且此时的树结构以 v 为根。那么接下来我们不断地进行换根的操作，即能在 O(N) 的时间内求出每个节点为根的答案，实现了时间复杂度的优化。

代码：

class Solution {
    int[] ans;
    int[] sz;
    int[] dp;
    List<List<Integer>> graph;

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        ans = new int[N];
        sz = new int[N];
        dp = new int[N];
        graph = new ArrayList<List<Integer>>();
        for (int i = 0; i < N; ++i) {
            graph.add(new ArrayList<Integer>());
        }
        for (int[] edge: edges) {
            int u = edge[0], v = edge[1];
            graph.get(u).add(v);
            graph.get(v).add(u);
        }
        dfs(0, -1);
        dfs2(0, -1);
        return ans;
    }

    public void dfs(int u, int f) {
        sz[u] = 1;
        dp[u] = 0;
        for (int v: graph.get(u)) {
            if (v == f) {
                continue;
            }
            dfs(v, u);
            dp[u] += dp[v] + sz[v];
            sz[u] += sz[v];
        }
    }

    public void dfs2(int u, int f) {
        ans[u] = dp[u];
        for (int v: graph.get(u)) {
            if (v == f) {
                continue;
            }
            int pu = dp[u], pv = dp[v];
            int su = sz[u], sv = sz[v];

            dp[u] -= dp[v] + sz[v];
            sz[u] -= sz[v];
            dp[v] += dp[u] + sz[u];
            sz[v] += sz[u];

            dfs2(v, u);

            dp[u] = pu;
            dp[v] = pv;
            sz[u] = su;
            sz[v] = sv;
        }
    }
}