树中距离之和

难度: Hard

给定一个无向、连通的树。树中有 n 个标记为 0...n-1 的节点以及 n-1 条边。

给定整数 n 和数组 edges ， edges[i] = [a_i, b_i]表示树中的节点 a_i 和 b_i 之间有一条边。

返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。

示例 1:

输入: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释: 树如图所示。
我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) 
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。

示例 2:

输入: n = 1, edges = []
输出: [0]

示例 3:

输入: n = 2, edges = [[1,0]]
输出: [1,1]

提示:

1 <= n <= 3 * 10⁴
edges.length == n - 1
edges[i].length == 2
0 <= a_i, b_i < n
a_i != b_i
给定的输入保证为有效的树

Submission

运行时间: 175 ms

内存: 40.3 MB

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        g = [[] for _ in range(n)]
        for x, y in edges:
            g[x].append(y)
            g[y].append(x)

        ans = [0] * n
        size = [1] * n
        def dfs(x: int, fa : int, depth: int) -> None:
            ans[0] += depth
            for y in g[x]:
                if y != fa:
                    dfs(y, x, depth + 1)
                    size[x] += size[y]
        dfs(0, -1, 0)

        def reroot(x: int, fa: int) -> None:
            for y in g[x]:
                if y != fa:
                    ans[y] = ans[x] + n - 2 * size[y]
                    reroot(y, x)
        reroot(0, -1)
        return ans

Explain

该题解采用两遍DFS（深度优先搜索）的方法。首先，建立邻接表表示树。第一次DFS计算出根节点（假设为节点0）到所有其他节点的距离总和并存储在ans[0]中，同时计算每个节点为根的子树的大小存储在size数组中。第二次DFS则利用已知的根节点的答案和子树大小，通过换根DP（动态规划）的思想，推导出其他节点的答案。具体地，当从节点x到其子节点y时，ans[y] = ans[x] + n - 2 * size[y]，这是基于树中每当根节点改变时，到达某个节点的距离改变的性质计算的。这种方法有效地避免了重复计算每个节点到所有其他节点的距离，大大提高了效率。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        g = [[] for _ in range(n)]  # 创建邻接表
        for x, y in edges:
            g[x].append(y)  # 无向图双向连接
            g[y].append(x)

        ans = [0] * n  # 存储结果
        size = [1] * n  # 子树大小初始化
        def dfs(x: int, fa: int, depth: int) -> None:
            ans[0] += depth  # 根到所有节点的距离和
            for y in g[x]:
                if y != fa:
                    dfs(y, x, depth + 1)
                    size[x] += size[y]  # 计算子树大小
        dfs(0, -1, 0)  # 从根节点开始DFS

        def reroot(x: int, fa: int) -> None:
            for y in g[x]:
                if y != fa:
                    ans[y] = ans[x] + n - 2 * size[y]  # 利用父节点结果计算子节点结果
                    reroot(y, x)  # 递归到下一层
        reroot(0, -1)  # 从根节点开始第二次DFS
        return ans  # 返回结果数组

Explore

选择节点0作为初始的根节点主要是出于方便和约定的原因，因为在大多数编程问题中，数组或列表的索引通常从0开始。实际上，选择任何一个节点作为根节点进行算法的第一次DFS都是可以的，这不会影响最终结果的正确性。算法的核心在于正确计算每个节点作为根时的子树大小和距离和。由于树是一个无环连通图，因此从任何一个节点开始，都能遍历整棵树，并得出正确的子树大小和初次计算的距离和。

`size`数组用于存储每个节点的子树中节点的总数（包括该节点本身）。在第一次DFS中，每遍历到一个节点，就更新其父节点的`size`值，增加当前节点子树的大小。这个信息对于第二次DFS中距离的快速计算至关重要。了解每个节点的子树大小可以帮助我们快速计算当根节点从一个节点变更到其子节点时，其他所有节点到新根节点的距离之和。具体来说，当根节点变更时，原根节点子树中的节点到新根的距离都会增加，而新根节点子树外的节点到新根的距离则会减少，这种变化可以通过子树大小直接计算得出，而不需要逐个遍历。

这个公式的推导基于树的性质，当根节点从x变更到其子节点y时，对于y的子树中的所有节点，它们到y的距离相比到x的距离减少了1。而对于不在y的子树中的节点，包括x，它们到y的距离相比到x的距离增加了1。因此，可以将这种变化量表示为：`总距离的变化量 = -1 * (y子树中的节点数) + 1 * (不在y子树中的节点数)`。由于整棵树有n个节点，不在y子树中的节点数即为`n - size[y]`。因此，`总距离的变化量`可以进一步写成 `n - 2 * size[y]`。将这个变化量加到原根节点x的距离和上，即得到新根节点y的距离和：`ans[y] = ans[x] + n - 2 * size[y]`。