https://leetcode.cn/problems/sum-of-distances-in-tree/

树形DP

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

    public void dfs(int u,int f){
        dp[u]=0;
        sz[u]=1;
        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){
        answer[u]=dp[u];
        for(int v:graph.get(u)){
            if(v==f){
                continue;
            }
            else{
                int du=dp[u];
                int su=sz[u];
                int dv=dp[v];
                int sv=sz[v];

                dp[u]-=dp[v]+sz[v];
                sz[u]-=sz[v];
                dp[v]+=dp[u]+sz[u];
                sz[v]+=sz[u];
                //不断深搜直到最深,没有子树为止,然后保存dp值,就是结果answer
                dfs2(v,u);

                //这是为下一个循环服务的,应该恢复
                dp[u]=du;
                dp[v]=dv;
                sz[u]=su;
                sz[v]=sv;
            }
        }
    }

    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        answer = new int[n];
        sz = new int[n];
        dp = new int[n];
        graph = new ArrayList<List<Integer>>();
        if(n==1){
            answer[0] = 0;
            return answer;
        }

        if(n==2){
            answer[0] = 1;
            answer[1] = 1;
            return answer;
        }

        for(int i=0;i<n;i++){
            graph.add(new ArrayList<Integer>());
        }

        for(int[] edge: edges){
            int u=edge[0];
            int v=edge[1];
            graph.get(u).add(v);
            graph.get(v).add(u);
        }

        dfs(0,-1);
        dfs2(0,-1);

        return answer;
    }
}

解释如下

photo