834. Sum of Distances in Tree

https://leetcode.com/problems/sum-of-distances-in-tree/

• 9/5更新解答解析

题目解释

有一颗由n个节点(node)组成的无向树，节点标示着0 到 n-1 的数字，且这棵树的节点只会有一个边(edge)连接。

我们要回传每个点到其他点的距离(边)之总和

Example1

节点0的距离总和是 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) --> 1 + 1 + 2 + 2 + 2 = 8
节点1的距离总和是 dist(1,0) + dist(1,2) + dist(1,3) + dist(1,4) + dist(1,5) --> 1 + 2 + 3 + 3 + 3 = 12

解法:

首先，先上程序码吧!!!

程序码:

class Solution(object):
    def sumOfDistancesInTree(self, N, edges):
        graph = collections.defaultdict(set)
        for u, v in edges:
            # 把每个有互相连接的nodes存起来
            graph[u].add(v)
            graph[v].add(u)
            
        count = [1] * N
        ans = [0] * N
        def dfs(node = 0, parent = None):
            for child in graph[node]:
                if child != parent:
                    dfs(child, node)
                    count[node] += count[child] #子树有几个节点
                    ans[node] += ans[child] + count[child]
        
        
        
        def dfs2(node = 0, parent = None):
            for child in graph[node]:
                if child != parent:
                    ans[child] = ans[node] - count[child] + N - count[child]
                    dfs2(child, node)

        dfs()
        dfs2()
        return ans

这个解法有两个function，我们一个个来看

dfs():

首先，dfs() 有两个任务:

1. 第一个是计算每个节点的子树有几个节点

   例如范例1: 节点2的子树有4个节点
   

2. 第二个是计算根(节点0)的距离总和，而其他节点则记下各自的子节点数量

   为什麽要记下各自的子节点数量呢?

   因为节点的距离总和计算方式是这样: ans[i] = 独立子树的根的距离总和 + 独立子树的子节点数量

   以节点1为例: ans[1] = 0 + 0
   节点2为例: ans[2] = (0 + 1) + (0 + 1) + (0 + 1)
   节点0就是: ans[0] = (0 + 1) + (3 + 4) = 8

到这里为止其实已经算是解完这题了，只要用个回圈让dfs()轮流跑过每个节点并记录答案就可以回传了
但是这样会超过系统限制的时间(Time Limit Exceeded)，所以才需要dfs2()

dfs2():

这个function是用来计算每个节点的距离总和

而funciton里面的下面这条公式让我一直想不明白，所以昨天就决定先卡个位然後去睡觉了

ans[child] = ans[node] - count[child] + N - count[child]

今天起床後再看一次才终於知道它在干嘛

这个公式要拆成两半来看会比较清楚

我们以 ans[1] = ans[0] - count[1] + N - count[1] 来举例

1. 前半段的 ans[0] - count[1] 代表和节点1相比，有些点离节点0更远而离节点1更近，因此用节点0的距离总和，减掉这些距离节点1更近的节点数量，就能得到节点1距离总合；而这些离节点1更近的点，其实就是节点1的子节点数量。

2. 既然有些点离节点1更近，那其他节点就会反而离节点1更远吧，所以才要加上 N - count[1]

闲聊:

连续两天都是hard难度的题目让我有点吃不消啊

过去做hard题目都觉得很困难，即使有答案可以看还是需要多想一下，所以都会偷懒避开不写

不过，都参加铁人赛了就硬干下去吧 !!!