codeforces1118F2 Tree Cutting (Hard Version)

題面

題意

給出一棵一共有k種顏色的樹，每個點要麼沒有顏色，要麼是k種顏色中的一種顏色，每種顏色至少出現一次，現在要你刪去k-1條邊，使每個連通塊中有且僅有一種顏色，問有幾種刪法。

做法

首先判斷無解的情況，可以發現如果一個有顏色的點爲根的子樹不包括所有這個顏色的點，則這個點與其父親之間的邊不可以被刪去，我們可以據此將其父親也染爲與它顏色相同的顏色。
因此，我們可以進行bfs，如果這一層中有多個顏色t，則將這層中所有顏色爲t的點的父親的顏色也染爲t（若其父親已經有顏色了，且不爲t，則無解），如此每種顏色的點都連成了一個連通塊，然後我們可以不斷將沒有顏色的葉子節點去掉，使圖更加精簡。
接下來我們進行dp，dp[i][0]表示此時以最高點爲i的連通塊中沒有顏色，dp[i][1]表示此時以最高點爲i的連通塊中有某個顏色，可以發現點u與其兒子v是否連邊可以由上述狀態決定。
下面考慮狀態轉移：
1.若點u沒有顏色，則：
$dp[u][0]=\prod_{i\in son(u)}{dp[i][0]+dp[i][1]}$
它可以由所有狀態轉移過來。
$dp[u][1]=\sum_{i\in son(u)}{dp[i][1]*\prod_{j\in son(u),j \not=i}{dp[j][0]+dp[j][1]}}$
枚舉它的顏色由哪個兒子轉移而來
2.若點u由顏色，則：
$dp[u][0]=0$它所在的連通塊不可能沒有顏色
$dp[u][1]=\prod_{i\in son(u)}{dp[i][0]+dp[i][1]}$
同樣可以由所有顏色轉移而來

代碼

#include<bits/stdc++.h>
#define ll long long
#define N 300100
#define M 998244353
using namespace std;

ll n,m,bb,tt,ans=1,num[N],fa[N],bfn[N],cnt[N],ds[N],dp[N][2];
bool gg[N];
vector<ll>to[N];
queue<ll>que;

inline ll po(ll u,ll v)
{
    ll res=1;
    for(;v;)
    {
	if(v&1) res=res*u%M;
	u=u*u%M;
	v>>=1;
    }
    return res;
}

void dfs(ll now,ll last)
{
    ll i,t;
    if(num[now]) dp[now][1]=1;
    else dp[now][0]=1;
    for(i=0;i<to[now].size();i++)
    {
	t=to[now][i];
	if(t==last) continue;
	dfs(t,now);
	if(num[now]) dp[now][1]=dp[now][1]*(dp[t][0]+dp[t][1])%M;
	else
	{
	    dp[now][0]=dp[now][0]*(dp[t][0]+dp[t][1])%M;
	    dp[now][1]=(dp[now][1]+dp[t][1]*po(dp[t][0]+dp[t][1],M-2)%M)%M;
	}
    }
    if(!num[now]) dp[now][1]=dp[now][1]*dp[now][0]%M;
}

int main()
{
    ll i,j,p,q,t;
    cin>>n>>m;
    for(i=1;i<=n;i++) scanf("%lld",&num[i]),cnt[num[i]]++;
    for(i=1;i<n;i++)
    {
	scanf("%lld%lld",&p,&q);
	to[p].push_back(q);
	to[q].push_back(p);
	ds[p]++,ds[q]++;
    }
    que.push(1);
    for(;!que.empty();)
    {
	q=que.front();
	que.pop();
	bfn[++tt]=q;
	for(i=0;i<to[q].size();i++)
	{
	    p=to[q][i];
	    if(p==fa[q]) continue;
	    fa[p]=q;
	    que.push(p);
	}
    }
    for(j=n;j>=1;j--)
    {
	i=bfn[j];
	if(!num[i]) continue;
	if(cnt[num[i]]<=1 || num[fa[i]]==num[i])
	{
	    cnt[num[i]]--;
	    continue;
	}
	if(num[fa[i]])
	{
	    puts("0");
	    return 0;
	}
	num[fa[i]]=num[i];
    }
    for(i=1;i<=n;i++)
    {
	if(ds[i]==1 && !num[i])
	{
	    que.push(i);
	}
    }
    for(;!que.empty();)
    {
	q=que.front();
	que.pop();
	gg[q]=1;
	for(i=0;i<to[q].size();i++)
	{
	    p=to[q][i];
	    if(num[p]) continue;
	    ds[p]--;
	    if(ds[p]==1) que.push(p);
	}
    }
    for(i=1;i<=n;i++) if(num[i]) break;
    dfs(i,-1);
    cout<<dp[i][1];
}