Queries to find sum of distance of a given node to every leaf node in a Weighted Tree | Java program for the above problem ; 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 ; System . out . println ( v ) ;
import java . util . * ; import java . lang . * ; class GFG { static class pair { int first , second ; pair ( int f , int s ) { this . first = f ; this . second = s ; } } static final int N = ( int ) 1e5 + 5 ; static ArrayList < ArrayList < pair > > v ; static int [ ] dp = new int [ N ] ; static int [ ] leaves = new int [ N ] ; static int n ; static void dfs ( int a , int par ) { int leaf = 1 ; for ( pair i : v . get ( a ) ) { if ( i . first == par ) continue ; leaf = 0 ; dfs ( i . first , a ) ; } if ( leaf == 1 ) { leaves [ a ] += 1 ; } else { for ( pair i : v . get ( a ) ) { if ( i . first == par ) continue ; leaves [ a ] += leaves [ i . first ] ; dp [ a ] = dp [ a ] + dp [ i . first ] + leaves [ i . first ] * i . second ; } } } static void dfs2 ( int a , int par ) { for ( pair i : v . get ( a ) ) { if ( i . first == par ) continue ; int leafOutside = leaves [ a ] - leaves [ i . first ] ; dp [ i . first ] += ( dp [ a ] - dp [ i . first ] ) ; dp [ i . first ] += i . second * ( leafOutside - leaves [ i . first ] ) ; leaves [ i . first ] += leafOutside ; dfs2 ( i . first , a ) ; } } static void answerQueries ( int [ ] queries ) { dfs ( 1 , 0 ) ; dfs2 ( 1 , 0 ) ; for ( int i = 0 ; i < queries . length ; i ++ ) { System . out . println ( dp [ queries [ i ] ] ) ; } } public static void main ( String [ ] args ) { n = 5 ; v = new ArrayList < > ( ) ; for ( int i = 0 ; i <= n ; i ++ ) v . add ( new ArrayList < pair > ( ) ) ; v . get ( 1 ) . add ( new pair ( 4 , 4 ) ) ; v . get ( 4 ) . add ( new pair ( 1 , 4 ) ) ; v . get ( 1 ) . add ( new pair ( 2 , 2 ) ) ; v . get ( 2 ) . add ( new pair ( 1 , 2 ) ) ; v . get ( 2 ) . add ( new pair ( 3 , 3 ) ) ; v . get ( 3 ) . add ( new pair ( 2 , 3 ) ) ; v . get ( 2 ) . add ( new pair ( 5 , 5 ) ) ; v . get ( 5 ) . add ( new pair ( 2 , 5 ) ) ; int [ ] queries = { 1 , 3 , 5 } ; answerQueries ( queries ) ; } }