Queries to find sum of distance of a given node to every leaf node in a Weighted Tree | C ++ 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 ) ;
#include <bits/stdc++.h> NEW_LINE using namespace std ; const int N = 1e5 + 5 ; vector < vector < pair < int , int > > > v ( N ) ; vector < int > dp ( N ) ; vector < int > leaves ( N ) ; int n ; void dfs ( int a , int par ) { bool leaf = 1 ; for ( auto & i : v [ a ] ) { if ( i . first == par ) continue ; leaf = 0 ; dfs ( i . first , a ) ; } if ( leaf == 1 ) { leaves [ a ] += 1 ; } else { for ( auto & i : v [ a ] ) { if ( i . first == par ) continue ; leaves [ a ] += leaves [ i . first ] ; dp [ a ] = dp [ a ] + dp [ i . first ] + leaves [ i . first ] * i . second ; } } } void dfs2 ( int a , int par ) { for ( auto & i : v [ 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 ) ; } } void answerQueries ( vector < int > queries ) { dfs ( 1 , 0 ) ; dfs2 ( 1 , 0 ) ; for ( int i = 0 ; i < queries . size ( ) ; i ++ ) { cout << dp [ queries [ i ] ] << endl ; } } int main ( ) { n = 5 ; v [ 1 ] . push_back ( make_pair ( 4 , 4 ) ) ; v [ 4 ] . push_back ( make_pair ( 1 , 4 ) ) ; v [ 1 ] . push_back ( make_pair ( 2 , 2 ) ) ; v [ 2 ] . push_back ( make_pair ( 1 , 2 ) ) ; v [ 2 ] . push_back ( make_pair ( 3 , 3 ) ) ; v [ 3 ] . push_back ( make_pair ( 2 , 3 ) ) ; v [ 2 ] . push_back ( make_pair ( 5 , 5 ) ) ; v [ 5 ] . push_back ( make_pair ( 2 , 5 ) ) ; vector < int > queries = { 1 , 3 , 5 } ; answerQueries ( queries ) ; }