Codeforces 1118 F2 Tree Cutting (Hard Version) —— 树形DP

This way

题意：

树上有一些0-k颜色的点，让你去掉k-1条边，使得每个连通块只包含相同颜色的点（0可以当成任何颜色），问你有多少种办法

题解：

网上一些博客都说不清楚感觉。。这道题也要自己意会一下。
首先对于每一种颜色，它都会构成一个连通块，这个连通块内是不能去边的，同样，如果两个颜色的连通块相交了，那么答案就是0.怎么判断相交？我看到一种很优秀的办法：dfs回溯的时候，如果这个点是连通块的顶端了，也就是加到它身上的cnt=这个颜色的数量，那么就代表这个连通块结束，将当前点的颜色和数量置为0。否则，如果它的儿子有值，它本身也有值，说明这两个连通块相交了，因为儿子没有被置为0说明连通块还存在着。那么如果当前点的颜色是0，并且儿子有颜色，那么将这个点的颜色置为儿子的颜色，代表这个点在连通块内。
现在用dp[i][0/1]表示：i点的子树是否有与i点相同颜色 的情况数（这个有点难以理解）
有颜色的条件：原本就有颜色或者一个无颜色的点在有颜色的连通块中。判断方法上面已经讲了。
那么状态转移方程分为两种：
1.当前点有颜色：
$dp[i][0]=0$
$dp[i][1]=\prod_{ne\in son[i]}(dp[ne][0]+dp[ne][1])$
也就是说当前点没有颜色的情况是不存在的，当前点有颜色的情况是所有儿子的积
2.当前点无颜色：
$dp[i][0]=\prod_{ne\in son[i]}(dp[ne][0]+dp[ne][1])$
$dp[i][1]=\sum_{ne\in son[i]}dp[ne][1]*\prod_{s\in i\&\&s!=ne}(dp[s][0]+dp[s][1])$
也就是说没有颜色的时候的情况是所有儿子的情况积
有颜色的时候，也就是与每一个儿子在同一个集合内时的情况乘上所有其他儿子的情况

#include<bits/stdc++.h>
using namespace std;
const int N=3e5+5;
#define ll long long
const ll mod=998244353;
ll dp[N][2],num[N],cnt[N],mul[N],c[N];
struct node
{
    int to,next;
}e[N*2];
int tot,head[N];
void add(int x,int y)
{
    e[tot].to=y;
    e[tot].next=head[x];
    head[x]=tot++;
}
int f=0;
int son[N],siz;
void dfs(int x,int fa)
{
    cnt[x]=c[x]>0;
    for(int i=head[x];~i;i=e[i].next)
    {
        int ne=e[i].to;
        if(ne==fa)continue;
        dfs(ne,x);
        if(f)
            return ;
        if(c[x]==0)
            c[x]=c[ne];
        else if(c[ne]&&c[x]!=c[ne])
        {
            f=1;
            return ;
        }
        cnt[x]+=cnt[ne];
    }
    mul[0]=1;
    siz=0;
    for(int i=head[x];~i;i=e[i].next)
    {
        if(e[i].to==fa)continue;
        son[++siz]=e[i].to;
        mul[siz]=mul[siz-1]*(dp[e[i].to][0]+dp[e[i].to][1])%mod;
    }
    if(c[x])dp[x][1]=mul[siz];
    else
    {
        dp[x][0]=mul[siz];
        ll ret=1;
        for(int i=siz;i;i--)
        {
            dp[x][1]+=dp[son[i]][1]*mul[i-1]%mod*ret%mod;
            ret=ret*(dp[son[i]][1]+dp[son[i]][0])%mod;
        }
    }
    if(num[c[x]]==cnt[x])
        cnt[x]=c[x]=0;
}
int main()
{
    memset(head,-1,sizeof(head));
    int n,k,x,y;
    scanf("%d%d",&n,&k);
    for(int i=1;i<=n;i++)scanf("%d",&c[i]),num[c[i]]++;
    for(int i=1;i<n;i++)scanf("%d%d",&x,&y),add(x,y),add(y,x);
    dfs(1,0);
    printf("%lld\n",f?0ll:dp[1][1]);
}