834. Sum of Distances in Tree

'''
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 answer of length n where answer[i] is the sum of the distances between the ith node in the tree and all other nodes.
'''
#Since it is tree we dont need visited array just one variable parent is enough
#Question follows root shifting technique
#sum[child]=sum[parent]-count[child]+n-count[child]
#It means when you shift root to a child then count[child] nodes move towards the root by '1' and n-count[child] moves away by '1'
class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        count=[1]*n  #We initialised all counts with 1 means counting itself as 1. We can remove here 1 and while caluculating count of every node there you can add 1
        ans=[0]*n #In this array we store sum of distances from node to node
        g=defaultdict(set)
        for u,v in edges:
            g[u].add(v)
            g[v].add(u)

        def dfs(node=0,parent=-1): #This traversal helps in building count array and also sum of all distances from 0 to every other node and every other node to root
            for child in g[node]:
                if child!=parent: #If its not visited
                    dfs(child,node)
                    count[node]+=count[child]
                    ans[node]+=ans[child]+count[child] #sum[p]=count[c]+sum[c] for every child. This you can generalise

        def dfs1(node=0,parent=-1): # This traversal helps in root shifting technique and caluculate all requires sums
            for c in g[node]:
                if c!=parent:
                    ans[c]=ans[node]-count[c]+n-count[c]
                    dfs1(c,node)

        dfs()
        dfs1()
        return ans