834. Sum of Distances in Tree

An undirected, connected tree with N nodes labelled 0...N-1 and N-1 edges are given.

The ith edge connects nodes edges[i][0] and edges[i][1] together.

Return a list ans, where ans[i] is the sum of the distances between node i 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: 
Here is a diagram of the given tree:
  0
 / \
1   2
   /|\
  3 4 5
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.

Note: 1 <= N <= 10000

思路：直接BFS会TLE，优化degree为1的Node也TLE，接着想优化其他degree的Node，到此有点卡住

答案是建树，然后求出每个Node有多个子Node，这样只要求出res[root]，求root的child的时候，就知道有count[child]个Node的距离拉近了，N-count[i]个Node的距离变远了1，这样只要以原始root为起点做一遍BFS/DFS就可以了

class Solution(object):
    def sumOfDistancesInTree(self, n, edges):
        """
        :type N: int
        :type edges: List[List[int]]
        :rtype: List[int]
        """
        from collections import defaultdict
        d = defaultdict(set)
        for s,t in edges: 
            d[s].add(t)
            d[t].add(s)
        
        count = [1]*n
        res = [0]*n
        
        # calculate from child to root
        def dfs(root, vis):
            vis.add(root)
            for i in d[root]:
                if i not in vis:
                    dfs(i, vis) 
                    # now count[i], res[i] is ready
                    # here res[i] means use i as root, sum of all children to i
                    # so, only res[init_root] is correct
                    count[root] += count[i]
                    res[root] += res[i] + count[i]
        
        # calculate from root to child
        def dfs2(root, vis):
            vis.add(root)
            for i in d[root]:
                if i not in vis:
                    res[i] = res[root] - count[i] + n-count[i]
                    dfs2(i, vis)
                    
        dfs(0, set())
        dfs2(0, set())
        return res
            
            
        
s=Solution()
print(s.sumOfDistancesInTree(n = 2, edges = [[0,1]]))
print(s.sumOfDistancesInTree(n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]))

834. Sum of Distances in Tree

猜你喜欢