tree reconstruct 834

834. Sum of Distances in Tree Hard

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree and all other nodes.

 Example 1:

Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: The tree is shown above.
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.
Hence, answer[0] = 8, and so on.

Example 2:

Input: n = 1, edges = []
Output: [0]

Example 3:

Input: n = 2, edges = [[1,0]]
Output: [1,1]

 Constraints:

• 1 <= n <= 3 * 104
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai, bi < n
• ai != bi
• The given input represents a valid tree.

class Solution {
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        //0.建图，这么做可以加速每个节点访问其子节点的速度
        List<Integer>[] graph = new List[n];
        for(int i = 0; i < n; i++) graph[i] = new ArrayList();
        for(int[] edge : edges) {
            int a = edge[0], b = edge[1];
            graph[a].add(b);
            graph[b].add(a);
        }
        //将0作为root节点，计算所有子树的size
        //1.define a root node  0
        //2.dfs calculate all the subtree size
        int[] size = new int[n];
        subtreeSize(0, graph, -1, size);
        //计算root节点到所有节点的距离和
        //3.dfs calculate the root node's distance to all the other nodes
        int[] result = new int[n];
        result[0] = rootDistance(0, graph, -1, size);
        //最后递归根据父节点计算每个节点到其他点的距离和
        //4.dfs calculate all the node's distance to the other nodes base on the formula
        calculateSubTree(0, graph, -1, size, result);
        return result;
    }
    //计算子书size
    private int subtreeSize(int curr, List<Integer>[] graph, int parent, int[] size){
        int sum = 0;
        for(int child : graph[curr]){
            if(child == parent) continue;
            sum += subtreeSize(child, graph, curr, size);
        }
        size[curr] = sum + 1;
        return size[curr];
    }
    //计算子树距离和
    private int rootDistance(int curr, List<Integer>[] graph, int parent, int[] size){
        int result = 0;
        for(int child : graph[curr]){
            if(child == parent) continue;
            //离所有child的距离，是child的所有距离和 + 1 + child的所有child个数（因为每增加一层，你离所有孙子的距离都得加一层）
            result += rootDistance(child, graph, curr, size) + size[child];
        }
        return result;
    }
    //基于父节点计算子树距离和
    private void calculateSubTree(int curr, List<Integer>[] graph, int parent, int[] size, int[] result){
        if(parent != -1) result[curr] = result[parent] + graph.length - 2 * size[curr];
        for(int child : graph[curr]){
            if(child == parent) continue;
            calculateSubTree(child, graph, curr, size, result);
        }
    }
}