動機

看解答推導看不懂，看lee215一下就懂了，真是神!!

另外這題很漂亮、很美!!

Problem

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.

 

Example 1:

Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]Output: [8,12,6,10,10,10]Explanation: The tree is shown above.We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)equals 1 + 1 + 2 + 2 + 2 = 8.Hence, answer[0] = 8, and so on.

Example 2:

Input: n = 1, edges = []Output: [0]

Example 3:

Input: n = 2, edges = [[1,0]]Output: [1,1]

 

Constraints:

  • 1 <= n <= 3 * 104
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • ai != bi
  • The given input represents a valid tree.

Sol

求根的距離和很簡單 ret[now] += count[x]+ret[x]

count[x]是x為根的node總個數

所以上面的ret的意思是新加一個點就會延長每個點的距離1步

這樣可以求出ret[root]，那其他點呢? 可以看成root移動，每次移動一格，就會有靠近的與遠離的 也就是，有node的距離變短，有的變長，所以 ret[x] = ret[parent] - count[x] + (N-count[x])

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        count = [1] * n
        ret = [0] * n
        gh = defaultdict(list)
        for a,b in edges:
            gh[a].append(b)
            gh[b].append(a)
        def post(now, prev):
            for x in gh[now]:
                if x != prev:
                    post(x,now)
                    count[now] += count[x]
                    ret[now] += count[x]+ret[x]
        def pre(now, prev):
            for x in gh[now]:
                if x != prev:
                    ret[x] = ret[now] - count[x] + n - count[x]
                    pre(x,now)
        post(0,None)
        # ret[0] is correct, other will not include parent part!!
        pre(0,None) # fix ret[i] from ret[0]
        return ret