leetcode 834. 树中距离之和

题目描述

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

给定整数 n 和数组 edges ， edges[i] = [ai, bi]表示树中的节点 ai 和 bi 之间有一条边。

返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。
示例 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:

输入: n = 1, edges = []
输出: [0]
示例 3:

输入: n = 2, edges = [[1,0]]
输出: [1,1]
提示:
1 <= n <= 3 * 10⁴
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
ai != bi
给定的输入保证为有效的树

分析

首先考虑以u为根的树，u到其它节点距离之和如何求出。j是u的孩子节点，dp[u]表示u到以u为根的子树中其它节点的距离之和。则dp[u] = sum(dp[j] + cnt[j]);其中cnt[j]表示以j为根的子树中节点的总数。j有cnt[j] - 1个后代节点，u到这些节点的距离仅比j到它们的距离多1，一共多了cnt[j] - 1，再加上u到j的距离1，就是dp[j] + cnt[j]了。所以求dp[u]，只需要将u所有孩子节点的dp和cnt值求出来即可。
整理下思路就是：

int dfs1(int u,int fa) {
        cnt[u] = 1;
        int res = 0;
        for(int i = h[u];~i;i = ne[i]) {
            int j = e[i];
            if(j == fa) continue;
            int m = dfs1(j,u);
            cnt[u] += m;
            res += m + dp[j];
        }
        dp[u] = res;
        return cnt[u];
    }

求u到以u为根的子树中其它节点的距离之和的问题顶多算得上medium难度，本题之所以难度是hard是因为需要求出树中所有节点到其他节点的距离之和，如果对每个节点都调用下dfs，时间复杂度是平方级别的，而本题的数据规模有三万，建立无向图的话规模就是六万，平方级别的算法显然会TLE。
做一次dfs已经求出了u到所有节点的距离之和，对于以u为根的子树中任一个节点v，dp[v]表示v到其后代节点的距离之和，这是我们已经求出的信息。设ans[u]表示u到树中其它节点的距离之和，显然ans[u] = dp[u]，但是树中其它节点仅求出了到其后代的距离之和，还需要加上到非后代节点的距离之和。

还是以这张图为例，对根节点0做次dfs求出了0到其它节点的距离之和，同时像dp[2]表示2到其后代节点3 4 5的距离之和。显然最终的ans[2]就等于dp[2]再加上2到0和1的距离，设0到除了2这个分支以外节点的距离之和是s，那么2到2的非后代节点（0和1）的距离之和就是s + 2，这相当于把0看作2的孩子节点，和之前的dfs一样，加上孩子节点的dp数和cnt数，不过之前求出的dp[0]或者叫ans[0]统计了0到所有节点的距离和，所以要减去0到2这个分支的距离和才是0到1的距离和，即s = ans[0] - dp[2] - cnt[2]，所以ans[2] = dp[2] + s + n - cnt[2] = dp[2] + ans[0] - dp[2] -cnt[2] + n - cnt[2]，其中n - cnt[2]表示2的非后代节点的个数。
看起来推理挺麻烦的，其实很简单，2到所有节点的距离就是2到后代节点的距离加上到非后代节点的距离，非后代节点的距离可以通过2的父节点到所有节点的距离减去到2这个分支上所有节点的距离，再加上非后代节点的个数求出。
总结下本题的思路，首先对根节点做次dfs，求出每个节点到其后代节点的距离之和存到dp数组里。然后再对根节点做次dfs，根据父节点的ans值和自身的dp值自上而下求出每个节点的ans值。

代码

class Solution {
public:
    static const int N = 60005;
    int idx = 0,e[N],h[N],ne[N];
    int tot = 0,cnt[N];
    vector<int> ans,dp;
    void add(int a,int b) {
        e[idx] = b,ne[idx] = h[a], h[a] = idx++;
    }
    int dfs1(int u,int fa) {
        cnt[u] = 1;
        int res = 0;
        for(int i = h[u];~i;i = ne[i]) {
            int j = e[i];
            if(j == fa) continue;
            int m = dfs1(j,u);
            cnt[u] += m;
            res += m + dp[j];
        }
        dp[u] = res;
        return cnt[u];
    }
    void dfs2(int u,int fa) {//求出ans[u]，继而递归的求其孩子节点的ans值
        if(fa != -1)    ans[u] += ans[fa] - cnt[u] - dp[u] + tot - cnt[u];
        for(int i = h[u];~i;i = ne[i]) {
            int j = e[i];
            if(j != fa) {
                dfs2(j,u);
            }
        }
    }
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        memset(h,-1,sizeof h);
        tot = n;
        int u;
        for(auto t : edges) add(t[0],t[1]),add(t[1],t[0]);//建无向图，省得找根节点
        dp = vector<int>(n,0);
        dfs1(0,-1);
        ans = dp;//先加上到后代节点的距离
        dfs2(0,-1);
        return ans;
    }
};