0834：树中距离之和（2197 分）

2018-05-13 约 675 字预计阅读 2 分钟次阅读

力扣第 84 场周赛第 4 题

题目

给定一个无向、连通的树。树中有 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
给定的输入保证为有效的树

相似问题：

分析

典型的换根 dp，通用的方法是两次 dfs
- 第一次从下往上求出根节点 0 的值
- 第二次从上往下根据相邻节点的关系更新值
本题中，第一次求节点 u 到所有子节点的距离之和 f[u]
- 除了知道每个子结点 v 的值 f[v] 外，还需要知道 v 的节点个数 sz[v]
第二次，对于 u 的子节点 v
- v 的子节点到 v 的距离都比 u 少 1，而其它节点到 v 的距离都比 u 多 1
- 因此 f2[v]=f2[u]-sz[v]+n-sz[v]

解答

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32


# dfs序预处理树
def prep(edges):
    n = len(edges)+1
    g = [[] for _ in range(n)]
    for u,v in edges:
        g[u].append(v)
        g[v].append(u)
    fa = [-1]*n
    order = []
    sk = [0]
    while sk:
        u = sk.pop()
        order.append(u)
        g[u] = [v for v in g[u] if v!=fa[u]]
        for v in g[u]:
            fa[v] = u
            sk.append(v)
    return g,order,fa

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        g,order,_ = prep(edges)
        f = [0]*n
        sz = [1]*n
        for u in order[::-1]:
            for v in g[u]:
                sz[u] += sz[v]
                f[u] += f[v]+sz[v]
        for u in order:
            for v in g[u]:
                f[v] = f[u]+n-sz[v]-sz[v]
        return f

131 ms

目录

0834：树中距离之和（2197 分）

题目

分析

解答