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);
}
}
}