Sub | CPP code to find the sub - tree with minimum color difference in a 2 - coloured tree ; Tree traversal to compute minimum difference ; Initial min difference is the color of node ; Traversing its children ; Not traversing the parent ; If the child is adding positively to difference , we include it in the answer Otherwise , we leave the sub - tree and include 0 ( nothing ) in the answer ; DFS for colour difference : 1 colour - 2 colour ; Minimum colour difference is maximum answer value ; Clearing the current value to check for colour2 as well ; Interchanging the colours ; DFS for colour difference : 2 colour - 1 colour ; Checking if colour2 makes the minimum colour difference ; Driver code ; Nodes ; Adjacency list representation ; Edges ; Index represent the colour of that node There is no Node 0 , so we start from index 1 to N ; Printing the result
#include <bits/stdc++.h> NEW_LINE using namespace std ; void dfs ( int node , int parent , vector < int > tree [ ] , int colour [ ] , int answer [ ] ) { answer [ node ] = colour [ node ] ; for ( auto u : tree [ node ] ) { if ( u == parent ) continue ; dfs ( u , node , tree , colour , answer ) ; answer [ node ] += max ( answer [ u ] , 0 ) ; } } int maxDiff ( vector < int > tree [ ] , int colour [ ] , int N ) { int answer [ N + 1 ] ; memset ( answer , 0 , sizeof ( answer ) ) ; dfs ( 1 , 0 , tree , colour , answer ) ; int high = 0 ; for ( int i = 1 ; i <= N ; i ++ ) { high = max ( high , answer [ i ] ) ; answer [ i ] = 0 ; } for ( int i = 1 ; i <= N ; i ++ ) { if ( colour [ i ] == -1 ) colour [ i ] = 1 ; else colour [ i ] = -1 ; } dfs ( 1 , 0 , tree , colour , answer ) ; for ( int i = 1 ; i < N ; i ++ ) high = max ( high , answer [ i ] ) ; return high ; } int main ( ) { int N = 5 ; vector < int > tree [ N + 1 ] ; tree [ 1 ] . push_back ( 2 ) ; tree [ 2 ] . push_back ( 1 ) ; tree [ 1 ] . push_back ( 3 ) ; tree [ 3 ] . push_back ( 1 ) ; tree [ 2 ] . push_back ( 4 ) ; tree [ 4 ] . push_back ( 2 ) ; tree [ 3 ] . push_back ( 5 ) ; tree [ 5 ] . push_back ( 3 ) ; int colour [ ] = { 0 , 1 , 1 , -1 , -1 , 1 } ; cout << maxDiff ( tree , colour , N ) ; return 0 ; }