LeetCode 834. 树中距离之和

一个简单的树形dp。两边dfs，第一遍统计出所有子节点到当前节点的距离。然后第二遍dfs用父节点更新子节点

class Solution {
public:
    const static int MAXN = 10005 ;
    int dp[MAXN][2], son[MAXN];
    int from[MAXN*2], to[MAXN*2], next[MAXN*2], head[MAXN] , tot;
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        init() ;
        for( int i = 0; i < edges.size(); i++ ){
            int u = edges[i][0] ;
            int v = edges[i][1] ;
            add( u, v );
            add( v, u );
        }
        dfs1( 0, -1 );
        dfs2( 0, -1 , N );
        vector<int> ans ;
        for( int i = 0 ; i < N; i++ ){
            ans.push_back( dp[i][1] );
        }
        return ans ;
    }
    void init(){
        memset(head, -1, sizeof head );
        tot = 0 ;
    }

    void add( int u, int v ){
        from[tot] = u ;
        to[tot] = v ;
        next[tot] = head[u] ;
        head[u] = tot++ ;
    }

    void dfs1( int u, int f ){
        dp[u][0] = 0 ;
        son[u] = 1 ;
        for( int i = head[u] ; i != -1; i = next[i] ){
            int v = to[i] ;
            if( v == f ) continue ;
            dfs1( v, u );
            dp[u][0] += dp[v][0] + son[v] ;
            son[u] += son[v] ;
        }
    }

    void dfs2( int u, int f, int n ){
        if( f == -1 ){
            dp[u][1] = dp[u][0];
        }
        for( int i = head[u]; i != -1; i = next[i] ){
            int v = to[i];
            if( v == f ) continue ;
            dp[v][1] = dp[u][1] + ( n - 2 * son[v] ) ;
            dfs2( v, u, n ) ;
        }
    }
};