好题不问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 ;
}