20201006
题目:
给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
示例:
输入: 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,以此类推。
说明: 1 <= N <= 10000
抛砖引玉
思路
单看一个子树:
0
/ \
1 2
/|\
3 4 5
- 节点 2:answer[2] = 6 -> 4 + 2
- 4 个直接子节点 + 到 0 距离为 2 的子节点
可以发现 0 到 2 的关系边被计算两次,将 0-1 看做 2 的子树,则该子数的距离变成了:子节点数+子节点作为子树根节点的距离
0
/
1
子树 dp[0] = 1
1
子树 dp[1] = 0
已知子树距离和子树个数组合整树距离:
设任意节点(v)为跟节点,假设链接其的节点(u)为子节点:
- dp[u] = dp[u] – dp[v] – sz[v] ,减去 u 中链接 v 的距离
- sz[u] = su[u] – sz[v] , 更新 u 下子节点数量
将 u 节点的数据追加到 v 节点下:
- dp[v] = dp[v] + dp[u] + sz[u]
- sz[v] = sz[v] + sz[u]
/**
* @param {number} N
* @param {number[][]} edges
* @return {number[]}
*/
var sumOfDistancesInTree = function(N, edges) {
let ans = Array(N).fill(0),
sz = Array(N).fill(0), // 子树节点个数
dp = Array(N).fill(0), // 子树根节点到其他节点和距离和
graph = Array(N)
.fill(0)
.map((v) => [])
// 循环树构建关系(记录与指定节点相邻的节点)
for (let [u, v] of edges) {
graph[u].push(v)
graph[v].push(u)
}
dfs(0, -1)
dfs2(0, -1)
// 计算子树根节点到其他节点的距离和
function dfs(u, f) {
sz[u] = 1
dp[u] = 0
for (let v of graph[u]) {
// 排除父节点
if (v === f) continue
dfs(v, u)
dp[u] += dp[v] + sz[v]
sz[u] += sz[v]
}
}
// 组合子树根节点距离得到整树距离
function dfs2(u, f) {
ans[u] = dp[u]
for (let v of graph[u]) {
// 排除父节点
if (v === f) continue
let pu = dp[u],
pv = dp[v],
su = sz[u],
sv = sz[v]
// v作为跟节点
dp[u] = dp[u] - (dp[v] + sz[v])
sz[u] = sz[u] - sz[v]
dp[v] = dp[v] + dp[u] + sz[u]
sz[v] = sz[v] + sz[u]
dfs2(v, u)
// 恢复子树节点数和子树距离
dp[u] = pu
dp[v] = pv
sz[u] = su
sz[v] = sv
}
}
return ans
}
正文完