题解
题目大意,给一个树,每个节点可能是红色蓝色或无色,删去一条边,要求删去后得到的两个树只有红色或蓝色节点 问有多少种删法,树至少有一个红色和蓝色节点。
DFS整个树,回溯时统计以每个节点为根的子树蓝色和红色节点数量,如果子树拥有某个颜色全部节点且没另一个颜色节点则满足条件。
AC代码
#include <stdio.h>
#include <bits/stdc++.h>
#define fst first
#define sed second
using namespace std;
typedef long long ll;
const int INF = 0x3f3f3f3f;
const ll LINF = 0x3f3f3f3f3f3f3f3f;
const int N = 3e5 + 10;
int a[N], b[N], c[N]; //以儿子节点为根的子树红蓝节点数量
vector<int> e[N];
int ans;
void DFS(int x, int f)
{
if (a[x] == 1)
b[x]++;
if (a[x] == 2)
c[x]++;
for (int y : e[x]) if (y != f)
{
DFS(y, x);
b[x] += b[y];
c[x] += c[y];
if (b[y] == b[0] && !c[y] || c[y] == c[0] && !b[y]) //儿子节点只拥有某一种颜色的全部节点且没有另一种颜色
ans++;
}
}
int main()
{
#ifdef LOCAL
freopen("C:/input.txt", "r", stdin);
#endif
int n;
cin >> n;
for (int i = 1; i <= n; i++)
{
scanf("%d", &a[i]);
if (a[i] == 1) //统计两个颜色节点总数
b[0]++;
if (a[i] == 2)
c[0]++;
}
for (int i = 1; i < n; i++)
{
int u, v;
scanf("%d%d", &u, &v);
e[u].push_back(v);
e[v].push_back(u);
}
DFS(1, 0);
cout << ans << endl;
return 0;
}