LeetCode 834. Sum of Distances in Tree

题目链接：https://leetcode.com/problems/sum-of-distances-in-tree/

题意：给出一棵树，对于每个节点，求出其他所有节点与该节点的距离之和。节点数不超过10000。

思路：如图所示，暴力复杂度太高，考虑进行优化，如图所示，考虑一对父子节点u与v，son[u]表示以u为根节点的子树的节点个数。若已知父节点u到其他所有节点的距离和，则可推出子节点到其他所有节点的距离和：以子节点为根的树上的点到该子节点的距离是到其父节点的距离-1，而其他节点到该子节点的距离是到其父节点的距离+1（图中黄色和红色部分）。即$ans[v]= ans[u]-son[v]+cnt-son[v]$，其中$cnt$为节点总数。每个节点的子节点数可通过一次dfs求出，求所有节点的距离和也可以一次dfs求出：先求父节点，再求子节点。

class Solution {
public:
    //dfs求出每个节点的子节点（包括自己）
    void dfs(int pre,int u, vector<int> g[],int height){
        int s0=0;
        son[u]=1;
        for(int i=0;i<g[u].size();i++){
            int v = g[u][i];
            if(v==pre)
                continue;
            ans[0] +=height+1;
            dfs(u,v,g,height+1);
            son[u]+=son[v];
        }
    }
    //dfs求出每个节点到其他所有节点的距离和
    void Dfs(int pre,int u,vector<int> g[]){
        for(int i=0;i<g[u].size();i++){
            int  v = g[u][i];
            if(v==pre)
                continue;
            ans[v]= ans[u]-son[v]+cnt-son[v];
            Dfs(u,v,g);
        }
    }
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        //总节点数
        cnt = N;
        vector<int> g[N];
        for(int i=0;i<edges.size();i++){
            g[edges[i][0]].push_back(edges[i][1]);
            g[edges[i][1]].push_back(edges[i][0]);
        }
        vector <int> s;
        if(N==1){
            s.push_back(0);
            return s;
        }
        son = new int[N];
        ans = new int[N];
        memset(son, 0, 4*N );
        ans[0]=0;
        dfs(-1,0, g ,0);
        s.push_back(ans[0]);
        Dfs(-1,0,g);
        for(int i=1;i<N;i++)
            s.push_back(ans[i]);
        return s;
    }
private:
    int *son = nullptr,*ans;
    int cnt;
};