2022-12-22.md

题目

834. Sum of Distances in Tree

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree and all other nodes.

思路

选择任意节点作为根节点 dfs1计算出所有节点的count和根节点的answer 注意dfs1中answer[node]表示node到所有子节点的距离和，而不是node到所有节点的距离和

count[node] = 1 + sum(count[child])
answer[node] = sum(answer[child] + count[child])

dfs2计算出所有节点的answer

answer[child] = answer[node] - count[child] + n - count[child]

代码

var sumOfDistancesInTree = function(n, edges) {
  const graph = {}
  for (let i = 0; i < n; i++) {
    graph[i] = []
  }
  for (const edge of edges) {
    graph[edge[0]].push(edge[1])
    graph[edge[1]].push(edge[0])
  }
  const count = []
  const answer = []
  function dfs1(node, parent) {
    count[node] = 1
    answer[node] = 0
    for (const child of graph[node]) {
      if (child == parent) continue
      dfs1(child, node)
      count[node] += count[child]
      answer[node] += answer[child] + count[child]
    }
  }
  function dfs2(node, parent) {
    for (const child of graph[node]) {
      if (child == parent) continue
      answer[child] = answer[node] - count[child] + n - count[child]
      dfs2(child, node)
    }
  }
  dfs1(0, 0)
  dfs2(0, 0)
  return answer
}