树中距离之和
题目原文
给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
第 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,以此类推。
题目解析
本题需要解决的问题有:
- 所给树是无向、连通树,只给定了树的边,因此树的根节点是不确定的,可能有多个,因此,不能像原本处理树时候,从根节点遍历。
- 假设给定一个根节点,我们可以容易计算根节点与其他节点的距离之和。但是对于其他节点,再用相同方法计算,计算复杂度过大,容易超出时间限制。我们可以改变原本树的根节点,并观察改变节点后,距离计算的改变。
解题方法
树形动态规划
定义 表示以u为根的子树,它的所有子节点到它的距离之和,同时定义 表示以u为根的子树的节点数量(包括u本身,u的孩子节点及孙子节点),不难得出如下的转移方程:
该公式可以理解为,对于节点u的距离和可以通过它的孩子节点的距离和计算而来,同时由于增加了u-v之间的边,影响了v和以v为根节点的子树的所有节点,因此,需要加上sons[v]。
我们遍历整棵树,从叶子节点开始自底向上递推到根节点,即深度优先搜索,dp[root],即为节点root的答案。
换根
对于本题来说,需要求解所有节点的距离和,从暴力的角度,我们可以对每个节点都做一次如上的树形动态规划,这样的时间复杂度为。
假设u
的某个孩子节点为v
,**让v换到根的位置,u变为v的节点,同时维护原有的dp信息。**可以观察到,这样的变化,不会影响u和v的其他节点dp值,因此只需要更新dp[u]和dp[v]的值。
观察dp转移方程,将u变为v的孩子节点后,对于u来说,需要减去v对于u的dp值的贡献。
注意:我们要同时更新u的sons值,sons[u]-sons[v]
同理,对于v来说,它的孩子节点增加了u,即:
我们也需要同时更新v的sons值。
解题代码
class Solution: def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]: if len(edges)==0: return [0] # 树形动态规划 result = [0 for _ in range(N)] # 该子树下所有节点,包括孩子和孙子节点 sons = [1 for _ in range(N)] tree = [[] for _ in range(N)] for parent,child in edges: tree[parent].append(child) tree[child].append(parent) # 先计算初始树的dp数组 dp = [0 for _ in range(N)] def dfs1(u,f): for v in tree[u]: if v == f: continue dfs1(v,u) dp[u] += dp[v]+sons[v] sons[u] += sons[v] # 从根节点开始 dfs1(0,-1) # 遍历其他节点,进行换根 def dfs2(u,f): result[u]=dp[u] # 假设将根节点u 换为当前节点v # 对于dp数组来说,仅需要改变u和v # 当u换为v的孩子节点时,需要去掉原来dp[v]对dp[u]的影响,因此dp[u]=dp[u]-(dp[v]+son[v]),同时,son[u]也要减去son[v] # 对于dp[v] dp[v]=dp[v]+dp[u]+son(u) son[v]也要加上son[u] for v in tree[u]: if v==f: continue pu = dp[u] pv = dp[v] su = sons[u] sv = sons[v] dp[u] -= dp[v] + sons[v] sons[u] -= sons[v] dp[v] += dp[u] + sons[u] sons[v] += sons[u] dfs2(v,u) dp[u] = pu dp[v] = pv sons[u] = su sons[v] = sv dfs2(0,-1) return result