Skip to content

Latest commit

 

History

History
50 lines (30 loc) · 1.26 KB

834.-sum-of-distances-in-tree.md

File metadata and controls

50 lines (30 loc) · 1.26 KB

834. Sum of Distances in Tree

class Solution:
    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:
        graph = collections.defaultdict(list)
        ## 关键是根是 所有子树的和的累加,
        ## 子树和就是每一个节点,做子树根时候,完美契合后序

        ## 其他的节点, 爸爸的结果,减去自身做根时候的个数,加上剩余节点的个数 
        ## 前序产生 

        self.dist = [0] * N
        self.count = [0] * N
        for x, y in edges:
            graph[x].append(y)
            graph[y].append(x)

        self.get_dis(0, -1,graph)
        self.get_count(0, -1, graph)

        return self.dist

    def get_dis(self,root,prev, graph):
        for i in graph[root]:
            if i == prev: continue
            self.get_dis(i, root, graph)
            self.dist[root] += self.dist[i] + self.count[i]
            self.count[root] += self.count[i]
        self.count[root] += 1
            

        
    def get_count(self, root, prev , graph):
        for i in graph[root]:
            if i == prev: continue
            self.dist[i] = self.dist[root] - self.count[i] + len(graph) - self.count[i]
            self.get_count(i, root,graph)