leetcode-834. 树中距离之和

2023-12-30 17:08

文章标签 leetcode 834 中距离

本文主要是介绍leetcode-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

解题思路

看了题解做的，树形dp

首先求解简单一些的问题，即给定根节点，求根节点到其他节点的距离和。以示例的树为例，先求0到其他节点的距离和，应该为：1的距离和+2的距离和+0的子树节点数量。1和2的距离和比较好理解，加上子树的节点数量是因为，从0走到子树的任意一个节点上，其距离都比从1或者2出发多了1，所以加上节点数量即可。

对每个节点保存一个tuple，分别是以该节点为根的树的距离和，以及以该节点为根的树的节点数量（包括节点本身），则上述的树可以写为：

   0/    \
1       2/ | \3   4   5
(0,1) (0,1) (0,1)

下一步写为：

	   0/    \1       2
(0,1)     (3,4) = ((0+1)*3, 1*3+1)/ | \3   4   5(0,1) (0,1) (0,1)

更新到根节点有：

	   0(8,6) = (0+3+1+4, 1+4+1)/    \1       2
(0,1)     (3,4)/ | \3   4   5(0,1) (0,1) (0,1)

以0为根节点，最终的答案就是8

这样计算出来的是以其中1个节点作为根节点的结果，以所有节点作为根节点，则可以计算出最终答案，这种方法的时间复杂度是 $o(n^2)$ 。求的时候用后序遍历即可。

观察0的子节点2可以发现，以2为根节点时，其实就是到2的子树的所有节点距离都缩短了1，到除了2的子树以外的节点都增加了1，所以ans[2] = ans[0] - num_child[2] + (N - num_child[2])。所以求其他节点时，不用重新计算，直接在已经计算出来的基础上变换即可。由于每次变换的基础都是父节点和子节点，所以用先序遍历，用父节点的值调整子节点的值即可。

最终的时间复杂度是 $o (n)$

注意点
由于上述的思路都是基于树的，而题目给出的是边，所以要从边中构造一棵树。因为给定的边的顺序不一致，所以先按照图的方式保存相连的节点，选定一个节点用BFS/DFS调整边的顺序，使得按照这种顺序从根节点开始画一棵树。

代码

class Solution:def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:if not edges:return [0]graph = {}for from_node, to_node in edges:if from_node not in graph:graph[from_node] = []graph[from_node].append(to_node)if to_node not in graph:graph[to_node] = []graph[to_node].append(from_node)# modify nodes order to form a treenode_dict = {}stack = [edges[0][0]]visited_nodes = set()while stack:node = stack.pop()if node in visited_nodes:continuevisited_nodes.add(node)if list(set(graph[node]) - visited_nodes):node_dict[node] = list(set(graph[node]) - visited_nodes)stack += node_dict.get(node, [])# postorderans = [0] * Nchild_num = [1] * Nstack = [(edges[0][0], 0)]while stack:node, stat = stack.pop()if stat == 0:stack.append((node, 1))if node in node_dict:for child_node in node_dict[node]:stack.append((child_node, 0))else:if node not in node_dict:continueans[node] = sum(ans[child_node] + child_num[child_node] for child_node in node_dict[node])child_num[node] = 1 + sum(child_num[child_node] for child_node in node_dict[node])# preorderstack = [edges[0][0]]while stack:node = stack.pop()if node in node_dict:for child_node in node_dict[node]:stack.append(child_node)ans[child_node] = ans[node] - child_num[child_node] + (N - child_num[child_node])return ans