leetcode t834-树中距离之和(树形DP)

 834. 树中距离之和

难度困难143收藏分享切换为英文接收动态反馈

给定一个无向、连通的树。树中有 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

设dp[cur]表示当前cur节点为根，所有子节点(无向树中，剩余节点皆是cur的子节点)到cur的距离和，sz[cur]表示cur所在子树的节点数目，所以,neb[cur]表示cur的子孙节点集合。那么简单的想，自底向上推一遍就可以得到单个节点的答案，所有节点i都做一遍，就可解。但是试了一下，这样做会超时，于是就照题解的思路，再做一遍dfs,以第一遍dfs的解集为基础，每访问一个当前cur的子孙节点，就将该子孙节点视为根，这样只需做如下处理:

，即旧根去掉新根的贡献，新根再补上旧根对它的贡献，更新答案之后将根再换回去。这样就用O(1)复杂度完成了一个节点答案的求解，整个问题总复杂度即O(n).

class Solution {
public:
    vector<int> res,dp,sz;
    vector<vector<int>> tree;
    void dfs(int cur, int pre){
        sz[cur]=1;
        dp[cur]=0;
        for(auto& v:tree[cur]){
            if(v==pre)continue;
            dfs(v,cur);
            dp[cur]+=dp[v]+sz[v];
            sz[cur]+=sz[v];
        }
    }
    void dfs2(int cur, int pre){
        res[cur]=dp[cur];
        for(auto& v:tree[cur]){
            if(v==pre)continue;
            int pu = dp[cur], pv = dp[v];
            int su = sz[cur], sv = sz[v];

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

            dfs2(v, cur);

            dp[cur] = pu, dp[v] = pv;
            sz[cur] = su, sz[v] = sv;
        }
    }
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        res.resize(N, 0);
        sz.resize(N, 0);
        dp.resize(N, 0);
        tree.resize(N, {});
        for(auto& edge:edges){
            tree[edge[0]].emplace_back(edge[1]);
            tree[edge[1]].emplace_back(edge[0]);
        }
        dfs(0,-1);
        dfs2(0,-1);
        return res;
    }
};