Tree Cutting (Easy Version) CodeForces - 1118F1(树形DP思想)

传送门

题意：给出一个树，然后对树上每个点进行染色，每次可以割去一条边，使得两边要么只有蓝色，要么只有红色，白色无所谓，就是不能一边既有红色又有蓝色，问图中有多少条这样的边？

题解：从1开始跑子树，维护当前点以及子树的红色个数和蓝色个数，之后再从1开始跑每条边，如果子树的红色个数等于所有红色个数并且蓝色个数为0，或者红色满蓝色空，那么以可以割去这条边，通过维护子树来进行。

附上代码：

#include<bits/stdc++.h>
using namespace std;
const int maxn=3e5+5;
int n,col[maxn],red[maxn],blue[maxn],re,bl,ans;
struct edge{int v,next;}e[maxn*2];
int head[maxn],tot;
void add_edge(int u,int v)
{
    e[++tot].v=v;e[tot].next=head[u];head[u]=tot;
    e[++tot].v=u;e[tot].next=head[v];head[v]=tot;
}
void dfs(int u,int fa)
{
    if(col[u]==1)red[u]++,re++;
    if(col[u]==2)blue[u]++,bl++;
    for(int i=head[u];i;i=e[i].next){
        if(e[i].v==fa)continue;
        dfs(e[i].v,u);
        red[u]+=red[e[i].v];
        blue[u]+=blue[e[i].v];
    }
}
void dfs1(int u,int fa)
{
    for(int i=head[u];i;i=e[i].next){
        if(e[i].v==fa)continue;
        if(blue[e[i].v]==bl&&red[e[i].v]==0||red[e[i].v]==re&&blue[e[i].v]==0)ans++;
        dfs1(e[i].v,u);
    }
}
int main()
{
    scanf("%d",&n);
    for(int i=1;i<=n;i++)scanf("%d",&col[i]);
    int u,v;
    for(int i=1;i<n;i++){
        scanf("%d%d",&u,&v);
        add_edge(u,v);
    }
    dfs(1,-1);
    dfs1(1,-1);
    printf("%d\n",ans);
    return 0;
}