给定一个由N个节点组成的N元树，以及一个矩阵N [] []，它由N – 1个形式为(X，Y)的边缘组成，表示节点X和节点Y之间的边缘，而数组col []由值组成：
- 0：无色节点。
- 1：节点显示为红色。
- 2：节点颜色为蓝色。
任务是查找仅由单色节点组成的给定树的子树数。
例子：
Input:
N = 5, col[] = {2, 0, 0, 1, 2},
edges[][] = {{1, 2}, {2, 3}, {2, 4}, {2, 5}}
Output: 1
Explanation:
A subtree of node 4 which is {4} has no blue vertex and contains only one red vertex.
Input:
N = 5, col[] = {1, 0, 0, 0, 2},
edges[][] = {{1, 2}, {2, 3}, {3, 4}, {4, 5}}
Output: 4
Explanation:
Below are the subtrees with the given property:
- Subtree with root node value 2 {2, 3, 4, 5}
- Subtree with root node value 3 {3, 4, 5}
- Subtree with root node value 4 {4, 5}
- Subtree with root node value 5 {5}
方法：可以使用深度优先搜索遍历来解决给定的问题。想法是使用给定树的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)