题目地址(834. 树中距离之和)

题目地址(834. 树中距离之和)

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

题目描述

给定一个无向、连通的树。树中有 n 个标记为 0...n-1 的节点以及 n-1 条边 。

给定整数 n 和数组 edges ， edges[i] = [ai, bi]表示树中的节点 ai 和 bi 之间有一条边。

返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。

 

示例 1:

输入: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释: 树如图所示。
我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) 
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。


示例 2:

输入: n = 1, edges = []
输出: [0]


示例 3:

输入: n = 2, edges = [[1,0]]
输出: [1,1]


 

提示:

1 <= n <= 3 * 104
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
ai != bi
给定的输入保证为有效的树

前置知识

公司

• 暂无

思路

关键点

代码

• 语言支持：Python3

Python3 Code:

class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        graph = [[] for _ in range(n)]
        for edge in edges:
            graph[edge[0]].append(edge[1])
            graph[edge[1]].append(edge[0])
        dist_sum = [0 for _ in range(n)]
        node_num = [1 for _ in range(n)]

        def post_order(node,parent):
            for i in graph[node]:
                if i == parent:
                    # 如果邻居只有父节点 则跳出循环
                    continue
                post_order(i,node)
                # node_num[node]记录node节点下的节点数量
                node_num[node] += node_num[i]
                # node_num[node]记录node子树下的路径距离和 计算局部dist_sum
                dist_sum[node] += dist_sum[i] + node_num[i]
        
        def pre_order(node,parent):
            for i in graph[node]:
                if i==parent:
                    continue
                # 调整根节点的位置 计算全局dist_sum
                dist_sum[i] = dist_sum[node] - node_num[i] + (n-node_num[i])
                pre_order(i,node)
        post_order(0,-1)
        pre_order(0,-1)
        return dist_sum    
        

复杂度分析

令 n 为数组长度。

• 时间复杂度：$O(n)$
• 空间复杂度：$O(n)$