[LeetCode] 834. Sum of Distances in Tree

LeetCode刷题记录

传送门

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 }