Count of subtrees from an N | C ++ program for the above approach ; Function to implement DFS traversal ; Mark node v as visited ; Traverse Adj_List of node v ; If current node is not visited ; DFS call for current node ; Count the total red and blue nodes of children of its subtree ; Count the no . of red and blue nodes in the subtree ; If subtree contains all red node & no blue node ; If subtree contains all blue node & no red node ; Function to count the number of nodes with red color ; Function to count the number of nodes with blue color ; Function to create a Tree with given vertices ; Traverse the edge [ ] array ; Create adjacency list ; Function to count the number of subtree with the given condition ; For creating adjacency list ; To store the count of subtree with only blue and red color ; visited array for DFS Traversal ; Count the number of red node in the tree ; Count the number of blue node in the tree ; Function Call to build tree ; DFS Traversal ; Print the final count ; Driver Code
#include <bits/stdc++.h> NEW_LINE using namespace std ; void Solution_dfs ( int v , int color [ ] , int red , int blue , int * sub_red , int * sub_blue , int * vis , map < int , vector < int > > & adj , int * ans ) { vis [ v ] = 1 ; for ( int i = 0 ; i < adj [ v ] . size ( ) ; i ++ ) { if ( vis [ adj [ v ] [ i ] ] == 0 ) { Solution_dfs ( adj [ v ] [ i ] , color , red , blue , sub_red , sub_blue , vis , adj , ans ) ; sub_red [ v ] += sub_red [ adj [ v ] [ i ] ] ; sub_blue [ v ] += sub_blue [ adj [ v ] [ i ] ] ; } } if ( color [ v ] == 1 ) { sub_red [ v ] ++ ; } if ( color [ v ] == 2 ) { sub_blue [ v ] ++ ; } if ( sub_red [ v ] == red && sub_blue [ v ] == 0 ) { ( * ans ) ++ ; } if ( sub_red [ v ] == 0 && sub_blue [ v ] == blue ) { ( * ans ) ++ ; } } int countRed ( int color [ ] , int n ) { int red = 0 ; for ( int i = 0 ; i < n ; i ++ ) { if ( color [ i ] == 1 ) red ++ ; } return red ; } int countBlue ( int color [ ] , int n ) { int blue = 0 ; for ( int i = 0 ; i < n ; i ++ ) { if ( color [ i ] == 2 ) blue ++ ; } return blue ; } void buildTree ( int edge [ ] [ 2 ] , map < int , vector < int > > & m , int n ) { int u , v , i ; for ( i = 0 ; i < n - 1 ; i ++ ) { u = edge [ i ] [ 0 ] - 1 ; v = edge [ i ] [ 1 ] - 1 ; m [ u ] . push_back ( v ) ; m [ v ] . push_back ( u ) ; } } void countSubtree ( int color [ ] , int n , int edge [ ] [ 2 ] ) { map < int , vector < int > > adj ; int ans = 0 ; int sub_red [ n + 3 ] = { 0 } ; int sub_blue [ n + 3 ] = { 0 } ; int vis [ n + 3 ] = { 0 } ; int red = countRed ( color , n ) ; int blue = countBlue ( color , n ) ; buildTree ( edge , adj , n ) ; Solution_dfs ( 0 , color , red , blue , sub_red , sub_blue , vis , adj , & ans ) ; cout << ans ; } int main ( ) { int N = 5 ; int color [ ] = { 1 , 0 , 0 , 0 , 2 } ; int edge [ ] [ 2 ] = { { 1 , 2 } , { 2 , 3 } , { 3 , 4 } , { 4 , 5 } } ; countSubtree ( color , N , edge ) ; return 0 ; }