树中的路径和 Sum of Distances in Tree

2019-03-28 15:25:43

问题描述：

问题求解：

写过的最好的Hard题之一。

初看本题，很经典的路径和嘛，dfs一遍肯定可以得到某个节点到其他所有节点的距离和。这种算法的时间复杂度是O(n ^ 2)。看一下数据量，emmm，果然不行。这个数据量一看就知道只能是O(n)的算法了。

只遍历一遍最多只能得到一个解，因此本题肯定是需要遍历至少两遍的。

在第一遍遍历的时候我们需要保存下两个值，一个是当前节点的subtree的路径总和，一个是当前节点的subtree的总的节点数。

在第二遍遍历的时候，我们已经知道root节点的值已经是最终的结果了，这个时候当我们将根节点从root移动到其child的时候，有cnt[child]的节点数的离现在的根节点近了一步，有N - cnt[child]的节点到当前根节点远了一步，所以res[child] = res[root] - cnt[child] + N - cnt[child]。这样再次遍历所有节点更新res数组得到的结果就是我们需要的最终的答案。

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        List<Set<Integer>> graph = new ArrayList<>();
        int[] res = new int[N];
        int[] cnt = new int[N];
        for (int i = 0; i < N; i++) graph.add(new HashSet<>());
        for (int[] edge : edges) {
            graph.get(edge[0]).add(edge[1]);
            graph.get(edge[1]).add(edge[0]);
        }
        dfs1(0, -1, res, cnt, graph);
        dfs2(0, -1, res, cnt, graph);
        return res;
    }

    private void dfs1(int root, int parent, int[] res, int[] cnt, List<Set<Integer>> graph) {
        for (int node : graph.get(root)) {
            if (node == parent) continue;
            dfs1(node, root, res, cnt, graph);
            cnt[root] += cnt[node];
            res[root] += res[node] + cnt[node];
        }
        cnt[root] += 1;
    }

    private void dfs2(int root, int parent, int[] res, int[] cnt, List<Set<Integer>> graph) {
        for (int node : graph.get(root)) {
            if (node == parent) continue;
            res[node] = res[root] - cnt[node]  + cnt.length - cnt[node];
            dfs2(node, root, res, cnt, graph);
        }
    }

2019-04-17 14:26:20

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        List<Set<Integer>> graph = new ArrayList<>();
        for (int i = 0; i < N; i++) graph.add(new HashSet<>());
        for (int[] edge : edges) {
            graph.get(edge[0]).add(edge[1]);
            graph.get(edge[1]).add(edge[0]);
        }
        int[] res = new int[N];
        int[] cnt = new int[N];
        int[] used = new int[N];
        used[0] = 1;
        dfs1(graph, 0, res, cnt, used);
        Arrays.fill(used, 0);
        used[0] = 1;
        dfs2(graph, 0, res, cnt, used);
        return res;
    }
    
    private int[] dfs1(List<Set<Integer>> graph, int root, int[] res, int[] cnt, int[] used) {
        res[root] = 0;
        cnt[root] = 1;
        for (int node : graph.get(root)) {
            if (used[node] == 1) continue;
            used[node] = 1;
            int[] r = dfs1(graph, node, res, cnt, used);
            res[root] += r[0] + r[1];
            cnt[root] += r[1];
        }
        return new int[]{res[root], cnt[root]};
    }
    
    private void dfs2(List<Set<Integer>> graph, int root, int[] res, int[] cnt, int[] used) {
        for (int node : graph.get(root)) {
            if (used[node] == 1) continue;
            used[node] = 1;
            res[node] = res[root] + cnt[0] - cnt[node]  * 2;
            dfs2(graph, node, res, cnt, used);
        }
    }