Description
An undirected, connected treewith 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
思路
題意：給出一棵樹，求出樹上每個節點到其他節點的距離總和。
題解：每個節點保存兩個值，一個是其子樹的節點個數（包括自身節點也要計數）nodesum[ ]，一個是其子樹各點到它的距離 dp[ ]，那么我們假設根節點為 u ，其僅有一個兒子 v ， u 到 v 的距離為 1 ，而 v 有若干兒子節點，那么 dp[v] 表示 v 的子樹各點到 v 的距離和，那么各個節點到達 u 的距離便可以這樣計算： dp[u] = dp[v] + nodesum[ v ] *1; （式子的理解，v 的一個兒子節點為 f，那么 f 到達 u 的距離為 (sum[ f ->v] + sum [v- > u])*1 ，dp[v] 包含了 sum[f->v]*1，所以也就是式子的分配式推廣到各個子節點計算出來的和）。我們已經知道了各個節點到達根節點的距離和，那么從根節點開始遞推下來可以得到各個點的距離和。另開一個數組表示每個節點的到其他節點的距離和，那么對於根節點u來說， dissum[u] = dp[u]。以 u 的兒子 v 為例， v 的子節點到 v 不必經過 v->u 這條路徑，因此 dissum[u] 多了 nodesum[v] * 1，但是對於不是 v 的子節點的節點，只到達了 u ，因此要到達 v 必須多走 u->v 這條路徑，因此 dissum[u] 少了 ( N - nodesum[v] ) * 1) ,所以 dissum[v] = dissum[u] - nodesum[v] * 1 + (N - nodesum[v] ) * 1，按照這個方法遞推下去就可以得到各個點的距離和。
1 class Solution {
2 private int tot = 0;
3 private Edge[] edge;
4 private int[] head;
5 private int[] dp;
6 private int[] nodesum;
7 private int[] dissum;
8 public int[] sumOfDistancesInTree(int N, int[][] edges) {
9 edge = new Edge[2 * N + 5];
10 head = new int[N + 5];
11 dp = new int[N + 5];
12 nodesum = new int[N + 5];
13 dissum = new int[N];
14 Arrays.fill(head,-1);
15 for (int i = 0;i < edges.length;i++){
16 int u = edges[i][0];
17 int v = edges[i][1];
18 addedge(u,v);
19 addedge(v,u);
20 }
21 dfs1(0,0);
22 dissum[0] = dp[0];
23 dfs2(0,0,N);
24 return dissum;
25 }
26
27 public void addedge(int u,int v){
28 edge[tot] = new Edge();
29 edge[tot].u = u;
30 edge[tot].v = v;
31 edge[tot].next = head[u];
32 head[u] = tot++;
33 }
34
35 public void dfs1(int u,int fa){
36 dp[u] = 0;
37 nodesum[u] = 1;
38 for (int i = head[u];i != -1;i = edge[i].next){
39 int v = edge[i].v;
40 if (v == fa) continue;
41 dfs1(v,u);
42 dp[u] += dp[v] + nodesum[v];
43 nodesum[u] += nodesum[v];
44 }
45 }
46
47 public void dfs2(int u,int fa,int sum){
48 for (int i = head[u];i != -1;i = edge[i].next){
49 int v = edge[i].v;
50 if (v == fa) continue;
51 dissum[v] = dissum[u] - nodesum[v] + sum - nodesum[v];
52 dfs2(v,u,sum);
53 }
54 }
55 class Edge{
56 int u,v,next;
57 }
58 }