连通块问题，我们考虑树形 DP。
设 \(f_{u,j}\) 表示在以 \(u\) 为根的子树内，选的颜色集合为 \(a_{u},j\)（两个颜色都必须选）且必须选点 \(u\) 的情况下的连通块个数。
特殊的，\(f_{u,a_{u}}\) 表示颜色只有 \(a_{u}\) 的情况数。
对于 \(v\in son_u\)，有：
若 \(a_{u}=a_{v}\)：
对于每一个颜色 \(i\)，有：
\[f_{u,i}=f_{v,i}\times f_{u,a_{u}}+f_{v,a_{v}}\times f_{u,i}+f_{u,i}\times f_{v,i}+f_{u,i} \]特别的，有：
\[f_{u,a_{u}}=f_{u,a_{u}}\times (f_{v,a_{v}}+1) \]否则：
\[f_{u,a_{v}}=f_{u,a_{v}}\times (f_{v,a_{v}}+f_{v,a_{u}}+1)+f_{u,a_{u}}\times (f_{v,a_{v}}+f_{v,a_{u}}) \]但是，这样子时间复杂度太大，接受不了。
我们发现上面对于的操作分别为线段树合并和线段树单点查询、修改，最后的减去对应区间修改，以此优化即可。
注意代码细节。
代码：
#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=1e5+5;
const ll mod=998244353;
int n,a[N];
struct node{
int l,r;
ll sum,mul;
}tree[N<<6];
int root[N],tot;
void clean(int k){
tree[k].sum%=mod,tree[k].mul%=mod;
}
int build(){
int k=++tot;
tree[k].mul=1;
return k;
}
void push_down(int k){
if(tree[k].mul==1) return;
if(tree[k].l){
tree[tree[k].l].mul*=tree[k].mul;
tree[tree[k].l].sum*=tree[k].mul;
clean(tree[k].l);
}
if(tree[k].r){
tree[tree[k].r].mul*=tree[k].mul;
tree[tree[k].r].sum*=tree[k].mul;
clean(tree[k].r);
}
tree[k].mul=1;
}
void push_up(int k){
tree[k].sum=tree[tree[k].l].sum+tree[tree[k].r].sum;
clean(k);
}
ll query(int k,int l,int r,int pos){
if(!k) return 0;
if(l==r) return tree[k].sum%mod;
push_down(k);
int mid=(l+r)>>1;
if(pos<=mid) return query(tree[k].l,l,mid,pos);
else return query(tree[k].r,mid+1,r,pos);
}
int insert(int k,int l,int r,int pos){
if(!k) k=build();
if(l==r){
tree[k].sum=1;
return k;
}
push_down(k);
int mid=(l+r)>>1;
if(pos<=mid) tree[k].l=insert(tree[k].l,l,mid,pos);
else tree[k].r=insert(tree[k].r,mid+1,r,pos);
push_up(k);
return k;
}
int merge(int x,int y,int l,int r,ll tu,ll tv,int cu){
if(!x&&!y) return 0;
if(!x){//f[u][l-r]都没有
tree[y].sum*=tu;
tree[y].mul*=tu;
clean(y);
return y;
}
if(!y){
tree[x].sum*=(tv+1ll);
tree[x].mul*=(tv+1ll);
clean(x);
return x;
}
if(l==r){
if(l==cu) tree[x].sum=tree[x].sum*(tv+1)%mod;
else tree[x].sum=tree[y].sum*tu%mod+tree[x].sum*tv%mod+tree[x].sum*tree[y].sum%mod+tree[x].sum;
tree[x].sum%=mod;
return x;
}
int mid=(l+r)>>1;
push_down(x),push_down(y);
tree[x].l=merge(tree[x].l,tree[y].l,l,mid,tu,tv,cu);
tree[x].r=merge(tree[x].r,tree[y].r,mid+1,r,tu,tv,cu);
push_up(x);
return x;
}
int mul(int k,int l,int r,int pos,ll tu,ll tv,ll fv_cu){
if(!k) k=build();
if(l==r){
tree[k].sum=tree[k].sum*(fv_cu+tv+1)%mod+tu*(fv_cu+tv)%mod;
tree[k].sum%=mod;
return k;
}
push_down(k);
int mid=(l+r)>>1;
if(pos<=mid) tree[k].l=mul(tree[k].l,l,mid,pos,tu,tv,fv_cu);
else tree[k].r=mul(tree[k].r,mid+1,r,pos,tu,tv,fv_cu);
push_up(k);
return k;
}
ll ans=0;
vector<int> G[N];
void dfs(int u,int fa){
root[u]=insert(root[u],1,n,a[u]);
ll tu,tv,fv_cu;
for(auto v : G[u]){
if(v==fa) continue;
dfs(v,u);
tu=query(root[u],1,n,a[u]);tv=query(root[v],1,n,a[v]);
if(a[v]!=a[u]){
fv_cu=query(root[v],1,n,a[u]);
root[u]=mul(root[u],1,n,a[v],tu,tv,fv_cu);
}
else{
root[u]=merge(root[u],root[v],1,n,tu,tv,a[u]);
}
}
ll sum=tree[root[u]].sum;
ans=(ans+sum)%mod;
}
int main(){
scanf("%d",&n);
for(int i=1;i<=n;i++)scanf("%d",&a[i]);
int x,y;
for(int i=1;i<n;i++){
scanf("%d%d",&x,&y);
G[x].push_back(y);
G[y].push_back(x);
}
dfs(1,0);
printf("%lld\n",ans);
return 0;
}