Queries to find sum of distance of a given node to every leaf node in a Weighted Tree | MAX size ; Graph with { destination , weight } ; For storing the sum for ith node ; Leaves in subtree of ith . ; dfs to find sum of distance of leaves in the subtree of a node ; flag is the node is leaf or not ; Skipping if parent ; Setting flag to false ; Doing dfs call ; Doing calculation in postorder . ; If the node is leaf then we just increment the no . of leaves under the subtree of a node ; Adding num of leaves ; Calculating answer for the sum in the subtree ; dfs function to find the sum of distance of leaves outside the subtree ; Number of leaves other than the leaves in the subtree of i ; Adding the contribution of leaves outside to the ith node ; Adding the leafs outside to ith node 's leaves. ; Calculating the sum of distance of leaves in the subtree of a node assuming the root of the tree is 1 ; Calculating the sum of distance of leaves outside the subtree of node assuming the root of the tree is 1 ; Answering the queries ; Driver Code ; 1 ( 4 ) / \ ( 2 ) / \ 4 2 ( 5 ) / \ ( 3 ) / \ 5 3 ; Initialising tree
N = 10 ** 5 + 5 NEW_LINE v = [ [ ] for i in range ( N ) ] NEW_LINE dp = [ 0 ] * N NEW_LINE leaves = [ 0 ] * ( N ) NEW_LINE n = 0 NEW_LINE def dfs ( a , par ) : NEW_LINE INDENT leaf = 1 NEW_LINE for i in v [ a ] : NEW_LINE INDENT if ( i [ 0 ] == par ) : NEW_LINE INDENT continue NEW_LINE DEDENT leaf = 0 NEW_LINE dfs ( i [ 0 ] , a ) NEW_LINE DEDENT if ( leaf == 1 ) : NEW_LINE INDENT leaves [ a ] += 1 NEW_LINE DEDENT else : NEW_LINE INDENT for i in v [ a ] : NEW_LINE INDENT if ( i [ 0 ] == par ) : NEW_LINE INDENT continue NEW_LINE DEDENT leaves [ a ] += leaves [ i [ 0 ] ] NEW_LINE dp [ a ] = ( dp [ a ] + dp [ i [ 0 ] ] + leaves [ i [ 0 ] ] * i [ 1 ] ) NEW_LINE DEDENT DEDENT DEDENT def dfs2 ( a , par ) : NEW_LINE INDENT for i in v [ a ] : NEW_LINE INDENT if ( i [ 0 ] == par ) : NEW_LINE INDENT continue NEW_LINE DEDENT leafOutside = leaves [ a ] - leaves [ i [ 0 ] ] NEW_LINE dp [ i [ 0 ] ] += ( dp [ a ] - dp [ i [ 0 ] ] ) NEW_LINE dp [ i [ 0 ] ] += i [ 1 ] * ( leafOutside - leaves [ i [ 0 ] ] ) NEW_LINE leaves [ i [ 0 ] ] += leafOutside NEW_LINE dfs2 ( i [ 0 ] , a ) NEW_LINE DEDENT DEDENT def answerQueries ( queries ) : NEW_LINE INDENT dfs ( 1 , 0 ) NEW_LINE dfs2 ( 1 , 0 ) NEW_LINE for i in range ( len ( queries ) ) : NEW_LINE INDENT print ( dp [ queries [ i ] ] ) NEW_LINE DEDENT DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT n = 5 NEW_LINE v [ 1 ] . append ( [ 4 , 4 ] ) NEW_LINE v [ 4 ] . append ( [ 1 , 4 ] ) NEW_LINE v [ 1 ] . append ( [ 2 , 2 ] ) NEW_LINE v [ 2 ] . append ( [ 1 , 2 ] ) NEW_LINE v [ 2 ] . append ( [ 3 , 3 ] ) NEW_LINE v [ 3 ] . append ( [ 2 , 3 ] ) NEW_LINE v [ 2 ] . append ( [ 5 , 5 ] ) NEW_LINE v [ 5 ] . append ( [ 2 , 5 ] ) NEW_LINE queries = [ 1 , 3 , 5 ] NEW_LINE answerQueries ( queries ) NEW_LINE DEDENT