树中每个节点到根的距离之和

树中每个节点到根的距离之和

问题:

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

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

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

来源:力扣(LeetCode)
链接:https://leetcode.cn/problems/sum-of-distances-in-tree
著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。

输入: 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,以此类推。

    def sumOfDistancesInTree(self, n: int, edges: list[list[int]]) -> list[int]:
        relationTree=[[] for _ in range(n)]
        sonNode=[0]*n
        dp=[0]*n
        result=[0]*n

#我认为这个算法的第一个难点是重建关系,因为是个无向树型图,所以[0,1]中0和1都可以看做是另一方的根。下面这个relationTree就是把这些关系重建。索引为根节点,值为子树。
        for f,s in edges:
            relationTree[f].append(s)
            relationTree[s].append(f)

        def dfs0(root:int, father:int):
            #sonNode[i]保存着i节点有多少个子节点
            sonNode[root]=1
            #dp[i]保存着i节点的子节点到i的举例和。计算过程必须用深度优先的方式从最下层的子树开始计算。
            dp[root]=0
            for item in relationTree[root]:
                if item==father:
                    continue
                dfs0(item,root)
                #dfs0函数返回代表当路径已经到了叶子节点,item就是叶子节点,root是叶子节点的父节点。
                dp[root]+=dp[item]+sonNode[item]
                sonNode[root]+=sonNode[item]
        dfs0(0,-1)
        #dfs1的作用是换根,将与原来的根相连接的节点变成新的根,具体的交换过程可以看力扣834号题目的解题部分的动画。
        def dfs1(root:int, father:int):
            #这个方法也是使用深度优先算法,因为不想对每个节点重新计算一次,本函数是通过使用dfs0的结果"sonNode"和"dp"来生成以其它节点为根时的结果。因为之前已经把0节点为根时的结果计算出来了,所以保存该值到result[root]中
            result[root]=dp[root]
            for item in relationTree[root]:
                if item==father:
                    continue
               #保存当前的值,因为修改过后还要恢复,因为要保持0为根时的所有数据,其它的值都是根据它生成的。
                pd,pn = dp[root],sonNode[root]
                sd,sn = dp[item],sonNode[item]
                dp[root] -= (dp[item]+sonNode[item])
                sonNode[root]-=sonNode[item]
                dp[item]+=dp[root]+sonNode[root]
                sonNode[item]=n
                dfs1(item,root)
                dp[root],sonNode[root] = pd,pn
                dp[item],sonNode[item] = sd,sn

        dfs1(0,-1)

        return result
Comments are closed.