834. Sum of Distances in Tree

                
class Solution {
    List[] graph;
    int[] count;
    int[] res;
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        graph = new List[n];
        count = new int[n];
        res = new int[n];
        for (int i = 0; i < n; i++) {
            graph[i] = new ArrayList<>();
        }
        for (int[] edge: edges) {
            graph[edge[0]].add(edge[1]);
            graph[edge[1]].add(edge[0]);
        }
        dfs(0, 0);
        dfs2(0, 0);
        return res;
    }
    
    private void dfs(int root, int prev) {
        for (int neigh : graph[root]) {
            if (neigh == prev) continue;
            dfs(neigh, root);
            res[root] += res[neigh] + count[neigh];
            count[root] += count[neigh];
        }
        count[root] += 1;
    }
    
    private void dfs2(int root, int prev) {
        for (int neigh : graph[root]) {
            if (neigh == prev) continue;
            res[neigh] = res[root] - count[neigh] + count.length - count[neigh];
            dfs2(neigh, root);
        }
    }
}