树中距离之和

题目链接

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

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

换根

对于本题来说，需要求解所有节点的距离和，从暴力的角度，我们可以对每个节点都做一次如上的树形动态规划，这样的时间复杂度为$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])$
注意：我们要同时更新u的sons值，sons[u]-sons[v]

同理，对于v来说，它的孩子节点增加了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/