leetcode每日一题—834.树中距离之和

题目：
给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

思路：
tree[i]:用来记录第i个结点的孩子节点，及其父节点
depth[i]:用来记录结点i的层次，根结点位于第0层
count[i]:用来记录结点i和 结点i的子结点 的总数
answer[0]:根节点到各结点的距离，即根结点所有子结点的层次之和，即depth列表中各元素之和
answer[child]=answer[father]+(N-count[child])-count[child]=answer[father]+N-2*count[child]
（即父结点的answer，加上少走的路，减去多走的路。在父结点的answer 基础上：少走的路，即当前结点 到 除 当前结点自身和当前结点的所有子节点 之外的所有结点 都少走了1，总共就是N-count[child]。在父结点的answer 基础上：多走的路，即当前结点的父节点 到 当前结点和当前结点的所有子结点 都多走了1，总共就是count[child]

解答：

class Solution:
    def sumOfDistancesInTree(self, N: int, edges: List[List[int]]) -> List[int]:
        tree = [[] for _ in range(N)]
        for father, child in edges:
            tree[father].append(child)
            tree[child].append(father)
        depth = [0 for _ in range(N)]
        count = [0 for _ in range(N)]

        def dfsForDepthAndCount(father, baba):
            count[father] = 1
            for child in tree[father]:
                if child != baba:
                    depth[child] = depth[father] + 1
                    dfsForDepthAndCount(child, father)
                    count[father] += count[child]


        dfsForDepthAndCount(0, -1)
        answer = [0 for _ in range(N)]
        answer[0] = sum(depth)


        def dfsForAnswer(father, baba):
            for child in tree[father]:
                if child != baba:
                    answer[child] = answer[father] + N - 2 * count[child]
                    dfsForAnswer(child, father)

        dfsForAnswer(0, -1)
        return answer