834.树中距离之和——hard

给定一个无向、连通的树。树中有 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，以此类推。
说明: 1 <= N <= 10000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree

思路：
首先容易想到:定义dp[u] 表示以 u 为根的子树，它的所有子节点到它的距离之和，同时定义size[u] 表示以 u 为根的子树的节点数量，不难得出如下的转移方程:

可得到树中任一节点v到子树所有节点的距离和。
题目要求得到任一节点v到树中所有节点的距离和，比较暴力的方法是枚举每个点作为根节点，然后计算，时间复杂度为O(n)，显然不能接受。
对于节点u和它的某子节点v来说，需要做一次节点变换，将v变成树的根，对于u来说，v不再是它的子节点，所以
dp[u] -= dp[v] + size[v];
size[u] -= size[v];
对于v来说，u节点变为它的子节点:
dp[v] += dp[u] +size[u];
size[v] += size[u];
这样就能以O(2n)=>O(n)的复杂度得出答案。
复杂度分析:
时间复杂度O(n)
空间复杂度O(n)（图使用邻接表保存，n个节点共n-1条边）

var sumOfDistancesInTree = function (N, edges) {
    let ans = new Array(N).fill(0);
    let size = new Array(N).fill(0);
    let dp = new Array(N).fill(0);
    let graph = new Array(N).fill(undefined).map(() => []);
    const dfs1 = (u, pre) => {
        size[u] = 1;
        dp[u] = 0;
        for (const v of graph[u]) {
            if (v === pre) {
                continue;
            }
            dfs1(v, u);
            dp[u] += dp[v] + size[v];
            size[u] += size[v];
        }
    };
    const dfs2 = (u, pre) => {
        ans[u] = dp[u];
        for (const v of graph[u]) {
            if (v === pre) {
                continue;
            }
            //将原值缓存
            const pu = dp[u], pv = dp[v];
            const su = size[u], sv = size[v];
            dp[u] -= dp[v] + size[v];
            size[u] -= size[v];
            dp[v] += dp[u] + size[u];
            size[v] += size[u];
            dfs2(v, u);
            //回溯
            dp[u] = pu;
            dp[v] = pv;
            size[u] = su;
            size[v] = sv;
        }
    };
    for (let edge of edges) {
        const [u, v] = edge;
        graph[u].push(v);
        graph[v].push(u);
    }
    dfs1(0, -1);
    dfs2(0, -1);
    return ans;
};