Java实现 LeetCode 834 树中距离之和（DFS+分析）

834. 树中距离之和

给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。

第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。

返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

示例 1:

输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]

输出: [8,12,6,10,10,10]

解释:

如下为给定的树的示意图：

0

/

1 2

/|

3 4 5

我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)

也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。

说明: 1 <= N <= 10000

PS：

这个题看了半个小时没找到思路，只能去看题解，原来我的菜

class Solution {

    //这个题我看了半天，怎么看怎么超时，原来是中间有规律，该说不说，这个规律真的难找

  int[] ans, count;

    List<Set<Integer>> graph;

    int N;

    public int[] sumOfDistancesInTree(int N, int[][] edges) {

        this.N = N;

        graph = new ArrayList<Set<Integer>>();

        ans = new int[N];

        count = new int[N];

        Arrays.fill(count, 1);

        for (int i = 0; i < N; ++i)

            graph.add(new HashSet<Integer>());

        for (int[] edge: edges) {

            graph.get(edge[0]).add(edge[1]);

            graph.get(edge[1]).add(edge[0]);

        }

        dfs(0, -1);

        dfs2(0, -1);

        return ans;

    }

    public void dfs(int node, int parent) {

        for (int child: graph.get(node))

            if (child != parent) {

                dfs(child, node);

                //count[node]是以node为根的节点个数

                count[node] += count[child];

                //ans[node]是所有结点道node的距离

                //就是ans[child]+child与node的距离

                ans[node] += ans[child] + count[child];

            }

    }

    public void dfs2(int node, int parent) {

        for (int child: graph.get(node))

            if (child != parent) {

                // ans[node] += ans[child] + count[child];

                //这是ans[child]的前半部分，

                //后半部分是，不是以child为根的节点个数

                //后半部分是，因为，我们可以发现，每一个count都是多加了自己本身的，

                //也就是我们前面的child是多加了自己本身的，所以不是child的根节点都没加上

                ans[child] = ans[node] - count[child] + N - count[child];

                dfs2(child, node);

            }

    }

}