【一天一大 lee】树中距离之和 (难度:困难) - Day20201006

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

抛砖引玉

思路

单看一个子树：

节点 2：answer[2] = 6 -> 4 + 2
4 个直接子节点 + 到 0 距离为 2 的子节点

可以发现 0 到 2 的关系边被计算两次，将 0-1 看做 2 的子树，则该子数的距离变成了：子节点数+子节点作为子树根节点的距离

   0
  /
 1

子树 dp[0] = 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
}