0834. 树中距离之和

ITCharge大约 3 分钟

0834. 树中距离之和

标签：树、深度优先搜索、图、动态规划
难度：困难

题目链接

0834. 树中距离之和 - 力扣

题目大意

描述：给定一个无向、连通的树。树中有 $n$ 个标记为 $0 \sim n - 1$ 的节点以及 $n - 1$ 条边。

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

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

说明：

$1 \le n \le 3 \times 10^4$ 。
$edges.length == n - 1$ 。
$edges[i].length == 2$ 。
$0 \le ai, bi < n$ 。
$ai \ne bi$ 。
给定的输入保证为有效的树。

示例：

示例 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 = 2, edges = [[1,0]]
输出: [1,1]

解题思路

思路 1：树形 DP + 二次遍历换根法

最容易想到的做法是：枚举 $n$ 个节点，以每个节点为根节点进行树形 DP。

对于节点 $u$ ，定义 $dp[u]$ 为：以节点 $u$ 为根节点的树，它的所有子节点到它的距离之和。

然后进行一轮深度优先搜索，在搜索的过程中得到以节点 $v$ 为根节点的树，节点 $v$ 与所有其他子节点之间的距离之和 $dp[v]$ 。还能得到子树的节点个数 $sizes[v]$ 。

对于节点 $v$ 来说，其对 $dp[u]$ 的贡献为：节点 $v$ 与所有其他子节点之间的距离之和，再加上需要经过 $u \rightarrow v$ 这条边的节点个数，即 $dp[v] + sizes[v]$ 。

可得到状态转移方程为： $dp[u] = \sum_{v \in graph[u]}(dp[v] + sizes[v])$ 。

这样，对于 $n$ 个节点来说，需要进行 $n$ 次树形 DP，这种做法的时间复杂度为 $O(n^2)$ ，而 $n$ 的范围为 $[1, 3 \times 10^4]$ ，这样做会导致超时，因此需要进行优化。

我们可以使用「二次遍历换根法」进行优化，从而在 $O(n)$ 的时间复杂度内解决这道题。

以编号为 $0$ 的节点为根节点，进行两次深度优先搜索。

第一次遍历：从编号为 $0$ 的根节点开始，自底向上地计算出节点 $0$ 到其他的距离之和，记录在 $ans[0]$ 中。并且统计出以子节点为根节点的子树节点个数 $sizes[v]$ 。
第二次遍历：从编号为 $0$ 的根节点开始，自顶向下地枚举每个点，计算出将每个点作为新的根节点时，其他节点到根节点的距离之和。如果当前节点为 $v$ ，其父节点为 $u$ ，则自顶向下计算出 $ans[u]$ 之后，我们将根节点从 $u$ 换为节点 $v$ ，子树上的点到新根节点的距离比原来都小了 $1$ ，非子树上剩下所有点到新根节点的距离比原来都大了 $1$ 。则可以据此计算出节点 $v$ 与其他节点的距离和为： $ans[v] = ans[u] + n - 2 \times sizes[u]$ 。

思路 1：代码

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        graph = [[] for _ in range(n)]

        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)


        ans = [0 for _ in range(n)]

        sizes = [1 for _ in range(n)]
        def dfs(u, fa, depth):
            ans[0] += depth
            for v in graph[u]:
                if v == fa:
                    continue
                dfs(v, u, depth + 1)
                sizes[u] += sizes[v]

        def reroot(u, fa):
            for v in graph[u]:
                if v == fa:
                    continue
                ans[v] = ans[u] + n - 2 * size[v]
                reroot(v, u)

        dfs(0, -1, 0)
        reroot(0, -1)
        return ans

思路 1：复杂度分析

时间复杂度： $O(n)$ ，其中 $n$ 为树的节点个数。
空间复杂度： $O(n)$ 。