Count of subtrees from an N | 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
def Solution_dfs ( v , color , red , blue , sub_red , sub_blue , vis , adj , ans ) : NEW_LINE INDENT vis [ v ] = 1 ; NEW_LINE for i in range ( len ( adj [ v ] ) ) : NEW_LINE INDENT if ( vis [ adj [ v ] [ i ] ] == 0 ) : NEW_LINE INDENT ans = Solution_dfs ( adj [ v ] [ i ] , color , red , blue , sub_red , sub_blue , vis , adj , ans ) ; NEW_LINE sub_red [ v ] += sub_red [ adj [ v ] [ i ] ] ; NEW_LINE sub_blue [ v ] += sub_blue [ adj [ v ] [ i ] ] ; NEW_LINE DEDENT DEDENT if ( color [ v ] == 1 ) : NEW_LINE INDENT sub_red [ v ] += 1 ; NEW_LINE DEDENT if ( color [ v ] == 2 ) : NEW_LINE INDENT sub_blue [ v ] += 1 ; NEW_LINE DEDENT if ( sub_red [ v ] == red and sub_blue [ v ] == 0 ) : NEW_LINE INDENT ( ans ) += 1 ; NEW_LINE DEDENT if ( sub_red [ v ] == 0 and sub_blue [ v ] == blue ) : NEW_LINE INDENT ( ans ) += 1 ; NEW_LINE DEDENT return ans NEW_LINE DEDENT def countRed ( color , n ) : NEW_LINE INDENT red = 0 ; NEW_LINE for i in range ( n ) : NEW_LINE INDENT if ( color [ i ] == 1 ) : NEW_LINE INDENT red += 1 ; NEW_LINE DEDENT DEDENT return red ; NEW_LINE DEDENT def countBlue ( color , n ) : NEW_LINE INDENT blue = 0 ; NEW_LINE for i in range ( n ) : NEW_LINE INDENT if ( color [ i ] == 2 ) : NEW_LINE INDENT blue += 1 NEW_LINE DEDENT DEDENT return blue ; NEW_LINE DEDENT def buildTree ( edge , m , n ) : NEW_LINE INDENT u , v , i = 0 , 0 , 0 NEW_LINE for i in range ( n - 1 ) : NEW_LINE INDENT u = edge [ i ] [ 0 ] - 1 ; NEW_LINE v = edge [ i ] [ 1 ] - 1 ; NEW_LINE if u not in m : NEW_LINE INDENT m [ u ] = [ ] NEW_LINE DEDENT if v not in m : NEW_LINE INDENT m [ v ] = [ ] NEW_LINE DEDENT m [ u ] . append ( v ) NEW_LINE m [ v ] . append ( u ) ; NEW_LINE DEDENT DEDENT def countSubtree ( color , n , edge ) : NEW_LINE INDENT adj = dict ( ) NEW_LINE ans = 0 ; NEW_LINE sub_red = [ 0 for i in range ( n + 3 ) ] NEW_LINE sub_blue = [ 0 for i in range ( n + 3 ) ] NEW_LINE vis = [ 0 for i in range ( n + 3 ) ] NEW_LINE red = countRed ( color , n ) ; NEW_LINE blue = countBlue ( color , n ) ; NEW_LINE buildTree ( edge , adj , n ) ; NEW_LINE ans = Solution_dfs ( 0 , color , red , blue , sub_red , sub_blue , vis , adj , ans ) ; NEW_LINE print ( ans , end = ' ' ) NEW_LINE DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT N = 5 ; NEW_LINE color = [ 1 , 0 , 0 , 0 , 2 ] NEW_LINE edge = [ [ 1 , 2 ] , [ 2 , 3 ] , [ 3 , 4 ] , [ 4 , 5 ] ] ; NEW_LINE countSubtree ( color , N , edge ) ; NEW_LINE DEDENT