无向连接的树(不一定是二叉树),求每个节点到其他节点的距离和。
返回一个数组,数组的第i个元素就是第i个节点到其他所有节点的距离之和。
思路:
涉及无向图的构造和遍历,树的前序后序遍历,问题拆分。
能想到的最笨的办法就是以每个节点为root进行DFS遍历,这样就能求到每个节点到其他所有节点的距离。但是这样要遍历n次。
如何能在O(n)时间内解决?
先看个例子,看节点0到其他所有节点的距离
0
/ \
1 2
/ \
3 4
0到1,2的距离都是1,1到3,4的距离是1,
而0到3,4的距离则在1到3,4的距离基础上加1,因为又深了一层,
现在要求的是距离之和,
那么0到3,4的距离之和 = 2(1到3,4的距离之和) + 1* 2(3,4是2个节点,每个节点加深了一层,距离+1)
那么0到所有节点的距离之和呢?
就 = 0到它连接的1,2的距离之和2 + 1到3,4距离和 + 3,4节点数
= 左子树的距离和(1到3,4的距离和+左子树的节点数(1,3,4共3个))
+ 右子树的距离和(2到子树null的距离和 + 右子树的节点数(只有2一个))
所以需要知道以每个节点为root的子树有多少个节点,
是一个从下到上遍历的过程,
所以用后序遍历,把结果保存在nodeNums数组。
遍历完之后,我们能得到节点0到其他所有节点的距离之和。
那其他节点怎么办,以1为例来说明。
之前我们求了以1为root的子树的节点数,是3个(1,3,4),
distance_sum[0] = 0到1 + 0到3 + 0到4 + 0到2
而distance_sum[1] = 1到1 + 1到3 + 1到4 + 1到2
可以看到1,3,4都在以1为root的子树中,这3个节点到1的距离比到0的距离少了1,
而这个子树之外的节点0和2,到1的距离比到0的距离多了1,
那么distance_sum[1] = distance_sum[0] - nodeNums[1]*1 + (n - nodeNums[1]) * 1,
这样一层一层往下走,就能得到所有节点的distance_sum[i],
这是一个从上到下遍历的过程,所以用到树的前序遍历。
无向图的遍历中需要有flag记录节点是不是已经被访问过,
而现在是树,不会有环的出现,但是因为边是双向的,只需要看当前节点是不是又回到上一节点即可。
class Solution {
ArrayList<Integer>[] graph;
int[] dis;
int[] nodeNums;
public int[] sumOfDistancesInTree(int n, int[][] edges) {
dis = new int[n];
nodeNums = new int[n];
graph = new ArrayList[n];
for(int i = 0; i < n; i++) graph[i] = new ArrayList<Integer>();
//make the graph
for(int[] edge : edges) {
graph[edge[0]].add(edge[1]);
graph[edge[1]].add(edge[0]);
}
postOrder(0, -1); //get node numbers
preOrder(0, -1, n); //get distance
return dis;
}
//获取以每个node为root的子树的节点数
void postOrder(int root, int pre) {
//left & right for binary tree
//这里不一定是二叉树,可能是多叉树
for(Integer node : graph[root]) {
if(node == pre) continue;
postOrder(node,root);
nodeNums[root] += nodeNums[node];
dis[root] += dis[node] + nodeNums[node];
}
nodeNums[root] ++; //节点数加上root自己
}
void preOrder(int root, int pre, int n) {
for(Integer node : graph[root]) {
if(node == pre) continue;
//dis[root]现在是正确的,而和root连接的点到root的距离要-1,
//以node为root的子树中的所有节点都需要距离-1,
//所以dis[node] = dis[root]-1*nodeNums[node]
//但同时,以node为root的子树之外的其他节点到node的距离要+1
//这个过程是自上而下的
dis[node] = dis[root] - nodeNums[node] + n - nodeNums[node];
preOrder(node, root, n);
}
}
}