手撕LeetCode #834——Python

834. 树中距离之和

给定一个无向、连通的树。树中有 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，以此类推。

解题思路：
看上去这道题目很复杂，似乎要依次去求解每一个结点和树中其他结点的距离，时间复杂度至少都是O(N^2)。但仔细分析，我们还是能够找到规律的。

首先我们先存储每个结点与它相连的结点记录：

tree_link = [[] for _ in range(N)]
for father, child in edges:
  tree_link[father].append(child)
  tree_link[child].append(father)

题目已经告诉我们会有多少个结点(N), 那么可以先初始化两个数组来表示每个结点的深度和它的子结点数(包含它自己)：

depth = [0 for _ in range(N)]
count = [1 for _ in range(N)]

接下来，就是递归地去更新这两个数组：

# grand表示父亲的父亲
def get_Depth_Count(father, grand):
  # 遍历与父亲结点相连的结点
  for child in three_link[father]:
    # 如果该结点不是父亲的父亲结点，则表示该结点不是叶结点
    if child != grand:
      # 子结点的深度肯定比父结点深， +1
      depth[child] = depth[father] + 1
      # 继续往下更新子结点的深度
      get_Depth_Count(child, father)
      # 遍历完子结点后，将子结点的数量加到父结点上
      count[father] += count[child]

最后，便是如何去具体地求解每一个结点与其他结点的距离，也就是我们要找的答案。我们先想一想，如果只是计算根结点与其他结点(也就是它的子结点)距离，结果就是将各个子结点的depth相加求和。好了，那我们再来看看其他结点的计算：

结点2的距离

我们看看在结点2处，它与根结点的距离是1，那它与子结点的距离就是根结点与2的子结点距离-1；同理，1结点与2结点的距离就是1结点与根结点的距离+1。那么规律就找到了：
answer[2] = answer[0] + 与2同层结点数(不包括2) - 2的子结点数(包括2)，因为我们已经知道总共N的结点，所以右边第二项可以改写为N-count[2]：
answer[2] = answer[0] + N -2* (2的子结点数(包括2))
根据树的子树递归性：
求解公式可以写为：
answer[child] = answer[father] + N -2* (child的子结点数(包括child))

answer = [0 for _ in range(N)]
answer[0] = sum(depth)
def dfs(father, grand):
  for child in tree_link[father]:
    if child != grand:
      answer[child] = answer[father] + n - 2*count[child]
      dfs(child, father)

dfs(0, -1)

解题思路参考：

「手画图解」树中距离之和 | 树形DP | 思路详解