好题不问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.①
(因为一个睿智错误多调了一个小时
1 #include<cstdio> 2 #include<vector> 3 #include<cstring> 4 #include<algorithm> 5 #define vint vector<int> 6 #define vit vector<int>::iterator 7 using namespace std; 8 const int N=300005,mod=998244353; 9 int p[N],noww[2*N],goal[2*N],siz[N]; 10 int anc[N],dep[N],imp[N],top[N],col[N],dp[N][2]; 11 int n,k,cnt,t1,t2; vint ve[N]; 12 void Link(int f,int t) 13 { 14 noww[++cnt]=p[f]; 15 goal[cnt]=t,p[f]=cnt; 16 noww[++cnt]=p[t]; 17 goal[cnt]=f,p[t]=cnt; 18 } 19 void Add(int &x,int y) 20 { 21 x+=y; 22 if(x>=mod) x-=mod; 23 } 24 void Mul(int &x,int y) 25 { 26 x=1ll*x*y%mod; 27 } 28 void DFS(int nde,int fth,int dth) 29 { 30 int tmp=0; 31 siz[nde]=1,anc[nde]=fth,dep[nde]=dth; 32 for(int i=p[nde];i;i=noww[i]) 33 if(goal[i]!=fth) 34 { 35 DFS(goal[i],nde,dth+1); 36 siz[nde]+=siz[goal[i]]; 37 if(siz[goal[i]]>tmp) 38 tmp=siz[goal[i]],imp[nde]=goal[i]; 39 } 40 } 41 void Decompose(int nde,int tpp) 42 { 43 top[nde]=tpp; 44 if(imp[nde]) 45 { 46 Decompose(imp[nde],tpp); 47 for(int i=p[nde];i;i=noww[i]) 48 if(goal[i]!=anc[nde]&&goal[i]!=imp[nde]) 49 Decompose(goal[i],goal[i]); 50 } 51 } 52 int LCA(int x,int y) 53 { 54 while(top[x]!=top[y]) 55 { 56 if(dep[top[x]]<dep[top[y]]) 57 swap(x,y); x=anc[top[x]]; 58 } 59 return dep[x]<dep[y]?x:y; 60 } 61 void Climb(int nde,int lca,int cor) 62 { 63 while(nde!=lca) 64 { 65 nde=anc[nde]; 66 if(col[nde]) 67 { 68 if(col[nde]==cor) return; 69 else printf("0"),exit(0); 70 } 71 col[nde]=cor; 72 } 73 } 74 void Getans(int nde,int fth) 75 { 76 vint v1,v2; 77 dp[nde][(bool)col[nde]]=1; 78 for(int i=p[nde];i;i=noww[i]) 79 if(goal[i]!=fth) 80 { 81 int g=goal[i]; Getans(g,nde); 82 int s=(dp[g][0]+dp[g][1])%mod; 83 v1.push_back(s),v2.push_back(s); 84 col[nde]?Mul(dp[nde][1],s):Mul(dp[nde][0],s); 85 } 86 if(!col[nde]&&v1.size()) 87 { 88 int sz=v1.size(),pt=0; 89 for(int i=1;i<sz;i++) Mul(v1[i],v1[i-1]); 90 for(int i=sz-2;i>=0;i--) Mul(v2[i],v2[i+1]); 91 for(int i=p[nde];i;i=noww[i]) 92 if(goal[i]!=fth) 93 { 94 int pre=pt?v1[pt-1]:1,suf=(pt==sz-1)?1:v2[pt+1]; 95 int tmp=dp[goal[i]][1]; 96 Mul(tmp,pre),Mul(tmp,suf),Add(dp[nde][1],tmp),pt++; 97 } 98 } 99 } 100 int main () 101 { 102 scanf("%d%d",&n,&k); 103 for(int i=1;i<=n;i++) 104 { 105 scanf("%d",&col[i]); 106 if(col[i]) ve[col[i]].push_back(i); 107 } 108 for(int i=1;i<n;i++) 109 scanf("%d%d",&t1,&t2),Link(t1,t2); 110 DFS(1,0,1),Decompose(1,1); 111 for(int i=1;i<=k;i++) 112 { 113 vint v=ve[i]; int lca=*v.begin(); 114 if(v.size()>1) 115 { 116 vit it=++v.begin(); 117 while(it!=v.end()) 118 lca=LCA(lca,*it++); 119 } 120 vit it=v.begin(); 121 while(it!=v.end()) 122 Climb(*it++,lca,i); 123 } 124 Getans(1,0); 125 // for(int i=1;i<=n;i++) 126 // printf("%d %d %d ",mar[i],dp[i][0],dp[i][1]); 127 printf("%d",dp[1][1]); 128 return 0; 129 }