834. Sum of Distances in Tree

,

DFS(TLE)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
class Solution {
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        boolean[][] edgeMap = new boolean[n][n];
        for (int[] edge : edges) {
            int from = edge[0], to = edge[1];
            edgeMap[from][to] = edgeMap[to][from] = true;
        }

        int[] res = new int[n];
        for (int i = 0; i < n; i++) {
            boolean[] visited = new boolean[n];
            res[i] = dfs(edgeMap,visited, i,0);
        }
        return res;
    }

    private int dfs(boolean[][] edgeMap, boolean[] visited, int start, int distance) {
        int distanceSum = distance;

        visited[start] = true;
        for (int i = 0; i < edgeMap[start].length; i++) {
            if (visited[i]) {
                continue;
            }
            if (edgeMap[start][i]) {
                distanceSum += dfs(edgeMap, visited, i, distance + 1);
            }
        }

        return distanceSum;
    }
}

DP

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
class Solution {
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        // 注意题目给出的 edges 是边的列表，并不是节点与节点集合的映射
        List<List<Integer>> nodeToNodesList = new ArrayList<>();
        for (int i = 0; i < n; i++) {
            nodeToNodesList.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            int from = edge[0], to = edge[1];
            nodeToNodesList.get(from).add(to);
            nodeToNodesList.get(to).add(from);
        }

        int[] size = new int[n];
        int[] answer = new int[n];

        // 在首次 dfs 的过程中计算出 0 节点与其他所有节点的距离之和，同时计算出所有节点的子树节点个数（含节点自身）
        dfs(answer, size, 0, -1, 0, nodeToNodesList);

        // 第二个 dfs, 计算出除 0 节点之外的节点与其他节点的距离之和
        distanceDfs(answer, size, 0, -1, nodeToNodesList);

        return answer;
    }

    private void distanceDfs(int[] answer, int[] size, int node, int parent, List<List<Integer>> nodeToNodesList) {
        if (parent != -1) {
            answer[node] = answer[parent] + (answer.length - size[node]) - size[node];
        }

        for (int nextNode : nodeToNodesList.get(node)) {
            if (nextNode != parent) {
                // 向下 dfs, 模拟二叉树遍历
                distanceDfs(answer, size, nextNode, node, nodeToNodesList);
            }
        }
    }

    private int dfs(int[] answer, int[] size, int node, int parent, int depth, List<List<Integer>> nodeToNodesList) {
        answer[0] += depth; // 加上当前节点与 0 节点的距离即可求出 0 节点与所有其他节点的距离之和
        int nodes = 1; // 当前节点算做一个

        // 遍历 node 的所有边
        for (int nextNode : nodeToNodesList.get(node)) {
            if (nextNode != parent) {
                // 向下 dfs, 模拟二叉树遍历
                nodes += dfs(answer, size, nextNode, node, depth + 1, nodeToNodesList);
            }
        }

        size[node] = nodes;
        return nodes;
    }
}

References

834. Sum of Distances in Tree