aboutsummaryrefslogtreecommitdiff
path: root/sum_dist_tree.py
blob: 258e96d1fab6678a5c5a0c60e94d5be7b14f490e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
# 834. Sum of Distances in Tree
from collections import defaultdict

"""
there is an undirected connected tree with 'n' nodes labeled from 0 to 'n -
1' and 'n - 1' edges. you are given the integer 'n' and the array 'edges'
where 'edges[i] = [ai, bi]' indicates that there is an edge between nodes
'ai' and 'bi' in the tree. return an array 'ans' of length 'n' where 'ans[i]'
is thesum of the distances between the i'th node in the tree and all other
nodes.
"""


class Solution(object):
    def sumOfDistancesInTree(self, n, edges):
        """
        :type n: int
        :type edges: List[List[int]]
        :rtype: List[int]
        """
        tree = defaultdict(set)
        res = [0] * n
        count = [1] * n
        for i, j in edges:
            tree[i].add(j)
            tree[j].add(i)

        def dfs(root, pre):
            for i in tree[root]:
                if i != pre:
                    dfs(i, root)
                    count[root] += count[i]
                    res[root] += res[i] + count[i]

        def dfs2(root, pre):
            for i in tree[root]:
                if i != pre:
                    res[i] = res[root] - count[i] + n - count[i]
                    dfs2(i, root)

        dfs(0, -1)
        dfs2(0, -1)
        return res


if __name__ == "__main__":
    obj = Solution()
    print(obj.sumOfDistancesInTree(n=6, edges=[[0, 1], [0, 2], [2, 3], [2, 4], [2, 5]]))
    print(obj.sumOfDistancesInTree(n=1, edges=[]))
    print(obj.sumOfDistancesInTree(n=2, edges=[[1, 0]]))