834. Sum of Distances in Tree

最新推荐文章于 2021-07-27 01:59:07 发布

BigFatSheep

最新推荐文章于 2021-07-27 01:59:07 发布

阅读量256

点赞数

CC 4.0 BY-SA版权

分类专栏： Leetcode Hard

本文链接：https://blog.youkuaiyun.com/BigFatSheep/article/details/84337135

Leetcode 同时被 2 个专栏收录

385 篇文章

订阅专栏

Hard

51 篇文章

订阅专栏

本文介绍了一种高效算法，用于解决给定树状结构中每个节点到所有其他节点距离之和的问题。通过两次深度优先搜索（DFS），算法能在O(N)时间内完成计算，避免了重复遍历。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Description

An undirected, connected tree with N nodes labelled 0…N-1 and N-1 edges are given.

The ith edge connects nodes edges[i][0] and edges[i][1] together.

Return a list ans, where ans[i] is the sum of the distances between node i 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:
Here is a diagram of the given tree:
0
/
1 2
/|
3 4 5
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.
Note: 1 <= N <= 10000

Problem URL

Solution

给一个二维数组表示的树，问以每一个节点为起点，到其他各个点的距离的和。

Intuition:
What if given a tree, with a certain root 0?
In O(N) we can find sum of distances in tree from root and all other nodes.
Now for all N nodes?
Of course, we can do it N times and solve it in O(N^2).
C++ and Java may get accepted luckly, but it’s not what we want.

When we move our root from one node to its connected node, one part of nodes get closer, one the other part get further.
If we know exactly hom many nodes in both parts, we can solve this problem.

With one single traversal in tree, we should get enough information for it and don’t need to do it again and again.

Explanation:
0. Let’s solve it with node 0 as root.

Initial an array of hashset tree, tree[i] contains all connected nodes to i.
Initial an array count, count[i] counts all nodes in the subtree i.
Initial an array of res, res[i] counts sum of distance in subtree i.

Post order dfs traversal, update count and res:
count[root] = sum(count[i]) + 1
res[root] = sum(res[i]) + sum(count[i])

Pre order dfs traversal, update res:
When we move our root from parent to its child i, count[i] points get 1 closer to root, n - count[i] nodes get 1 futhur to root.
res[i] = res[root] - count[i] + N - count[i]

return res, done.

Time Complexity:
dfs: O(N)
dfs2: O(N)

Code

class Solution {
    private int[] res;
    private int[] count;
    private ArrayList<HashSet<Integer>> tree;
    private int n;
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        tree = new ArrayList<>();
        res = new int[N];
        count = new int[N];
        n = N;
        for (int i = 0; i < N; i++){
            tree.add(new HashSet<Integer>());
        }
        for (int[] edge : edges){
            tree.get(edge[0]).add(edge[1]);
            tree.get(edge[1]).add(edge[0]);
        }
        dfs(0, new HashSet<Integer>());
        dfs2(0, new HashSet<Integer>());
        return res;
    }
    
    private void dfs(int root, HashSet<Integer> seen){
        seen.add(root);
        for (int i : tree.get(root)){
            if (!seen.contains(i)){
                dfs(i, seen);
                count[root] += count[i];
                res[root] += res[i] + count[i];
            }
        }
        count[root]++;
    }
    
    private void dfs2(int root, HashSet<Integer> seen){
        seen.add(root);
        for (int i : tree.get(root)){
            if (!seen.contains(i)){
                res[i] = res[root] - count[i] + n - count[i];
                dfs2(i, seen);
            }
        }
    }
}