求一颗树以任意点为根，其他点到根的和

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

 

 

 思路 ： 一遍 dfs 可以求到以任意一个点为根，所有点所含孩子的个数，cnt数组记录

　　　　考虑每一条边的贡献时，即此边左右节点数相乘

　　　　第2遍 dfs, 以其他点为根时，例如以2为根，先去掉 0-1 边的贡献，再加上以2为根 0和1 对 2 的贡献

代码示例 ： 

class Solution {
public:
    int cnt[10005];
    int ans[10005];
    vector<int>ve[10005];
    int pt[10005];

    void dfs(int x){
        //printf("++++%d  %d\n", x, ve[x].size());
        pt[x] = 1; cnt[x] = 1;
         
        int sum = 0;
        for(int i = 0; i < ve[x].size(); i++){
            if (!pt[ve[x][i]]) {
                dfs(ve[x][i]);
                sum += cnt[ve[x][i]];
                //printf("+++ %d %d %d %d \n", x, ve[x][i], cnt[x], cnt[ve[x][i]]);
                ans[0] += cnt[x] * cnt[ve[x][i]];
            }
        }
        cnt[x] += sum;
    }
    
    void dfs2(int x, int N){
        pt[x] = 1;    
        for(int i = 0; i < ve[x].size(); i++){
            if (!pt[ve[x][i]]) {
                ans[ve[x][i]] = (ans[x]-cnt[ve[x][i]])+(N-cnt[ve[x][i]]);
                dfs2(ve[x][i], N);
            }
        }
    }

    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        vector<int> sum;
        
        for(int i = 0; i < N-1; i++){
            int u = edges[i][0];
            int v = edges[i][1];
            ve[u].push_back(v);
            ve[v].push_back(u);
        }

        memset(cnt, 0, sizeof(cnt));
        memset(pt, 0, sizeof(pt));
        dfs(0);
        memset(pt, 0, sizeof(pt));
        dfs2(0, N);

        for(int i = 0; i < N; i++){
            sum.push_back(ans[i]);
        }

        return sum;
    }
};