树中距离之和

题目链接

题目原文 给定一个无向、连通的树。树中有 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,以此类推。

题目解析

本题需要解决的问题有:

  • 所给树是无向、连通树,只给定了树的边,因此树的根节点是不确定的,可能有多个,因此,不能像原本处理树时候,从根节点遍历。
  • 假设给定一个根节点,我们可以容易计算根节点与其他节点的距离之和。但是对于其他节点,再用相同方法计算,计算复杂度过大,容易超出时间限制。我们可以改变原本树的根节点,并观察改变节点后,距离计算的改变。

解题方法

树形动态规划

定义dp[u]dp[u] 表示以u为根的子树,它的所有子节点到它的距离之和,同时定义sons[u]sons[u] 表示以u为根的子树的节点数量(包括u本身,u的孩子节点及孙子节点),不难得出如下的转移方程:
dp[u]=vson[u](dp[v]+sons[v]) dp[u]=\sum_{v\in son[u]}(dp[v]+sons[v])
该公式可以理解为,对于节点u的距离和可以通过它的孩子节点的距离和计算而来,同时由于增加了u-v之间的边,影响了v和以v为根节点的子树的所有节点,因此,需要加上sons[v]。

我们遍历整棵树,从叶子节点开始自底向上递推到根节点,即深度优先搜索,dp[root],即为节点root的答案。

换根

对于本题来说,需要求解所有节点的距离和,从暴力的角度,我们可以对每个节点都做一次如上的树形动态规划,这样的时间复杂度为O(N)O(N)

假设u的某个孩子节点为v,**让v换到根的位置,u变为v的节点,同时维护原有的dp信息。**可以观察到,这样的变化,不会影响u和v的其他节点dp值,因此只需要更新dp[u]和dp[v]的值。

观察dp转移方程,将u变为v的孩子节点后,对于u来说,需要减去v对于u的dp值的贡献。
dp[u]=dp[u](dp[v]+sons[v])dp[u]=dp[u]-(dp[v]+sons[v])
注意:我们要同时更新u的sons值,sons[u]-sons[v]

同理,对于v来说,它的孩子节点增加了u,即:
dp[v]=dp[v]+(dp[u]+sons[u]) dp[v]=dp[v]+(dp[u]+sons[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

参考

https://leetcode-cn.com/problems/sum-of-distances-in-tree/solution/shu-zhong-ju-chi-zhi-he-by-leetcode-solution/