(127)834. 树中距离之和(leetcode)

题目链接：
https://leetcode-cn.com/problems/sum-of-distances-in-tree/
难度：困难
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，以此类推。
说明: 1 <= N <= 10000

不会。。。好tm难 官方题解
dp[u] 表示以 u 为根的子树，它的所有子节点到它的距离之和
sz[u] 表示以 u 为根的子树的节点数量
graph[u][i]=v 边 u->v(i=0,1,2…)

class Solution {
public:
    vector<int> ans, sz, dp;
    vector<vector<int>> graph;

    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
    	// 初始化
        ans.resize(N,0);
        sz.resize(N,0);
        dp.resize(N,0);
        graph.resize(N,{});
        for(auto& a:edges){
            int u=a[0],v=a[1];
            graph[u].emplace_back(v);
            graph[v].emplace_back(u);
        }
        // 求出dp[0]
        dfs1(0,-1);
        dfs2(0,-1);
        return ans;
    }

    void dfs1(int u,int f){
    	// 初始dp[u]为0 sz[u]为1（树 连通无环）
        dp[u]=0;
        sz[u]=1;
        // 对u的子节点进行枚举
        // dp[u]+=(dp[v]+sz[v])
        // dp[v]是 v 所有子节点到它的距离之和 sz[v]以 v 为根的子树的节点数量
        // dp[v]+sz[v] 节点v对dp[u]的贡献
        // sz[u]+=sz[v] v是u的子节点 sz[v]就是对u的子节点数量的贡献
        for(auto& v:graph[u]){
            if(v==f){
                continue;
            }
            dfs1(v,u);
            dp[u]+=dp[v]+sz[v];
            sz[u]+=sz[v];
        }
    }

    void dfs2(int u,int f){
        ans[u]=dp[u];
		// 将u的子节点作为根节点 并更新 dp[u] dp[v] sz[u] sz[v]
        for(auto& v:graph[u]){
            if(v==f){
                continue;
            }
            int dpu=dp[u],dpv=dp[v];
            int szu=sz[u],szv=sz[v];
            // 更新
            // dp[u]减去dp[v]的贡献  dp[v]加上dp[u]的贡献
            // sz[u]减去sz[v]的贡献  sz[v]加上sz[u]的贡献
            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]=dpu,dp[v]=dpv;
            sz[u]=szu,sz[v]=szv;
        }
    }
};