📌  相关文章
📜  由单色节点组成的 N 叉树的子树计数

📅  最后修改于: 2021-09-07 02:34:24             🧑  作者: Mango

给定一个由N 个节点组成的 N 元树和一个由 N – 1 个(X, Y)形式的边组成的矩阵边 [][]表示节点 X和节点Y之间的边,以及一个由值组成的数组 col[]

  • 0:无色节点。
  • 1:节点为红色。
  • 2:节点为蓝色。

任务是找到仅由单色节点组成的给定树的子树数量。
例子:

方法:给定的问题可以使用深度优先搜索遍历来解决。这个想法是使用给定树的 DFS 计算每个子树中红色和蓝色节点的数量。计算后,计算仅包含蓝色节点和仅包含红色节点的子树的数量。
下面是上述方法的实现:

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


Python3
# Function to implement DFS traversal
def Solution_dfs(v, color, red, blue, sub_red, sub_blue, vis, adj, ans):
  
    # Mark node v as visited
    vis[v] = 1;
  
    # Traverse Adj_List of node v
    for i in range(len(adj[v])):
  
        # If current node is not visited
        if (vis[adj[v][i]] == 0):
  
            # DFS call for current node
            ans=Solution_dfs(adj[v][i], color,red, blue,sub_red, sub_blue,vis, adj, ans);
  
            # Count the total red and blue
            # nodes of children of its subtree
            sub_red[v] += sub_red[adj[v][i]];
            sub_blue[v] += sub_blue[adj[v][i]];
         
    if (color[v] == 1):
        sub_red[v] += 1;
     
    # Count the no. of red and blue
    # nodes in the subtree
    if (color[v] == 2):
        sub_blue[v] += 1;
     
    # If subtree contains all
    # red node & no blue node
    if (sub_red[v] == red and sub_blue[v] == 0):
        (ans) += 1;
  
    # If subtree contains all
    # blue node & no red node
    if (sub_red[v] == 0 and sub_blue[v] == blue):
        (ans) += 1;
     
    return ans
      
# Function to count the number of
# nodes with red color
def countRed(color, n):
 
    red = 0;
     
    for i in range(n):
     
        if (color[i] == 1):
            red += 1;
     
    return red;
   
# Function to count the number of
# nodes with blue color
def countBlue(color, n):
 
    blue = 0;
     
    for i in range(n):
     
        if (color[i] == 2):
            blue += 1
     
    return blue;
   
# Function to create a Tree with
# given vertices
def buildTree(edge, m, n):
 
    u, v, i = 0, 0, 0
  
    # Traverse the edge[] array
    for i in range(n - 1):
  
        u = edge[i][0] - 1;
        v = edge[i][1] - 1;
  
        # Create adjacency list
        if u not in m:
            m[u] = []
             
        if v not in m:
            m[v] = []
             
        m[u].append(v)
        m[v].append(u);
         
# Function to count the number of
# subtree with the given condition
def countSubtree(color, n, edge):
  
    # For creating adjacency list
    adj = dict()
    ans = 0;
  
    # To store the count of subtree
    # with only blue and red color
    sub_red = [0 for i in range(n + 3)]
    sub_blue = [0 for i in range(n + 3)]
  
    # visted array for DFS Traversal
    vis = [0 for i in range(n + 3)]
  
    # Count the number of red
    # node in the tree
    red = countRed(color, n);
  
    # Count the number of blue
    # node in the tree
    blue = countBlue(color, n);
  
    # Function Call to build tree
    buildTree(edge, adj, n);
  
    # DFS Traversal
    ans=Solution_dfs(0, color, red, blue,sub_red, sub_blue, vis, adj, ans);
  
    # Print the final count
    print(ans, end = '')
     
# Driver Code
if __name__=='__main__':
 
    N = 5;
    color = [ 1, 0, 0, 0, 2 ]
    edge = [ [ 1, 2 ], [ 2, 3 ], [ 3, 4 ], [ 4, 5 ] ];
  
    countSubtree(color, N, edge);
   
# This code is contributed by rutvik_56


输出:
4

时间复杂度: O(N)
辅助空间: O(N)

如果您想与行业专家一起参加直播课程,请参阅Geeks Classes Live