解题：CF1118F2 Tree Cutting (Hard Version)

题面

好题不问Div(这是Div3最后一题，不得不说Mike真是强=。=)

首先同一个颜色的点的LCA要和它们在一个划分出的块里，那么我们先按颜色把所有点到它们的LCA的路径涂色，如果这个过程中出现了重合的颜色则说明无解。

之后问题转化为一个树形DP问题，设$dp[i][0/1]$表示以$i$为根的子树中i是否划入一个有颜色的块的方案数。然后讨论转移：

1.如果i自身没有颜色

①如果i不划入有颜色的块，那么直接继承子树信息，乘法原理统计，$dp[i][0]=\prod_{v∈son_i}(dp[v][0]+dp[v][1])$

②如果i划入有颜色的块，那么对于每一个子树都单独统计一遍划入其中的方案数，$dp[i][1]=\sum_{v∈son_i}dp[v][1]\prod_{w∈son_i\&\&w!=v}(dp[w][0]+dp[w][1])$，维护每个节点子节点的前缀后缀乘积来转移

2.如果i自身有颜色

①如果i不划入有颜色的块，那么......不划入有颜色的块=。=？？？$dp[i][0]=0$

②如果i划入有颜色的块，同理于1.①

(因为一个睿智错误多调了一个小时

 #include<cstdio>
 #include<vector>
 #include<cstring>
 #include<algorithm>
 #define vint vector<int>
 #define vit vector<int>::iterator
 using namespace std;
 const int N=,mod=;
 int p[N],noww[*N],goal[*N],siz[N];
 int anc[N],dep[N],imp[N],top[N],col[N],dp[N][];
 int n,k,cnt,t1,t2; vint ve[N];
 void Link(int f,int t)
 {
     noww[++cnt]=p[f];
     goal[cnt]=t,p[f]=cnt;
     noww[++cnt]=p[t];
     goal[cnt]=f,p[t]=cnt;
 }
 void Add(int &x,int y)
 {
     x+=y;
     if(x>=mod) x-=mod;
 }
 void Mul(int &x,int y)
 {
     x=1ll*x*y%mod;
 }
 void DFS(int nde,int fth,int dth)
 {
     int tmp=;
     siz[nde]=,anc[nde]=fth,dep[nde]=dth;
     for(int i=p[nde];i;i=noww[i])
         if(goal[i]!=fth)
         {
             DFS(goal[i],nde,dth+);
             siz[nde]+=siz[goal[i]];
             if(siz[goal[i]]>tmp)
                 tmp=siz[goal[i]],imp[nde]=goal[i];
         }
 }
 void Decompose(int nde,int tpp)
 {
     top[nde]=tpp;
     if(imp[nde])
     {
         Decompose(imp[nde],tpp);
         for(int i=p[nde];i;i=noww[i])
             if(goal[i]!=anc[nde]&&goal[i]!=imp[nde])
                 Decompose(goal[i],goal[i]);
     }
 }
 int LCA(int x,int y)
 {
     while(top[x]!=top[y])
     {
         if(dep[top[x]]<dep[top[y]])
             swap(x,y); x=anc[top[x]];
     }
     return dep[x]<dep[y]?x:y;
 }
 void Climb(int nde,int lca,int cor)
 {
     while(nde!=lca)
     {
         nde=anc[nde];
         if(col[nde])
         {
             if(col[nde]==cor) return;
             else printf(""),exit();
         }
         col[nde]=cor;
     }
 }
 void Getans(int nde,int fth)
 {
     vint v1,v2;
     dp[nde][(bool)col[nde]]=;
     for(int i=p[nde];i;i=noww[i])
         if(goal[i]!=fth)
         {
             int g=goal[i]; Getans(g,nde);
             int s=(dp[g][]+dp[g][])%mod;
             v1.push_back(s),v2.push_back(s);
             col[nde]?Mul(dp[nde][],s):Mul(dp[nde][],s);
         }
     if(!col[nde]&&v1.size())
     {
         int sz=v1.size(),pt=;
         for(int i=;i<sz;i++) Mul(v1[i],v1[i-]);
         for(int i=sz-;i>=;i--) Mul(v2[i],v2[i+]);
         for(int i=p[nde];i;i=noww[i])
             if(goal[i]!=fth)
             {
                 int pre=pt?v1[pt-]:,suf=(pt==sz-)?:v2[pt+];
                 int tmp=dp[goal[i]][];
                 Mul(tmp,pre),Mul(tmp,suf),Add(dp[nde][],tmp),pt++;
             }
     }
 }
 int main ()
 {
     scanf("%d%d",&n,&k);
     for(int i=;i<=n;i++)
     {
         scanf("%d",&col[i]);
         if(col[i]) ve[col[i]].push_back(i);
     }
     for(int i=;i<n;i++)
         scanf("%d%d",&t1,&t2),Link(t1,t2);
     DFS(,,),Decompose(,);
     for(int i=;i<=k;i++)
     {
         vint v=ve[i]; int lca=*v.begin();
         if(v.size()>)
         {
             vit it=++v.begin();
             while(it!=v.end())
                 lca=LCA(lca,*it++);
         }
         vit it=v.begin();
         while(it!=v.end())
             Climb(*it++,lca,i);
     }
     Getans(,);
 //    for(int i=1;i<=n;i++)
 //        printf("%d %d %d\n",mar[i],dp[i][0],dp[i][1]);
     printf("%d",dp[][]);
     return ;
 }