#include "joitour.h"
#define pb push_back
#define fi first
#define se second
#define mp make_pair
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
typedef long double ld;
template <typename T> bool chkmin(T &x,T y){return x>y?x=y,1:0;}
template <typename T> bool chkmax(T &x,T y){return x<y?x=y,1:0;}
int n,tot,ct;
int a[200005],v[400005],nxt[400005],h[200005],sz[200005],dfn[200005][20],dfo[200005][20],c[200005][3],cnt[200005][20][3],dep[200005],cf[200005][20],cr[200005][20];
bool vis[200005];
ll ans,cc[200005][3],ca[200005],pa[200005][20][3],sp[200005][3],pd[200005];
void addedge(int x,int y){
v[++tot]=y; nxt[tot]=h[x]; h[x]=tot;
v[++tot]=x; nxt[tot]=h[y]; h[y]=tot;
}
struct fenwick{
int val[8000005];
fenwick(){
for(int i=1;i<=8000000;i++) val[i]=0;
}
void change(int p,int v){
for(;p<=8000000;p+=p&-p) val[p]+=v;
}
int ask(int p){
int v=0;
for(;p;p-=p&-p) v+=val[p];
return v;
}
int ask(int l,int r){
return ask(r)-ask(l-1);
}
}sub[3],up[3];
void dfs1(int x,int f){
sz[x]=1;
for(int p=h[x];p;p=nxt[p]){
if(v[p]==f||vis[v[p]]) continue;
dfs1(v[p],x);
sz[x]+=sz[v[p]];
}
}
int dfs2(int x,int f,int k){
for(int p=h[x];p;p=nxt[p]){
if(vis[v[p]]||v[p]==f||sz[v[p]]*2<k) continue;
return dfs2(v[p],x,k);
}
return x;
}
void dfs3(int x,int f,int d,int rt,int r){
cf[x][d]=rt;
cr[x][d]=r;
c[rt][a[x]]++;
cnt[r][d][a[x]]++;
dfn[x][d]=++ct;
for(int p=h[x];p;p=nxt[p]){
if(v[p]==f||vis[v[p]]) continue;
dfs3(v[p],x,d,rt,r);
}
dfo[x][d]=++ct;
if(a[x]==1){
pa[r][d][0]+=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
sp[rt][0]+=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
pa[r][d][2]+=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
sp[rt][2]+=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
}
up[a[x]].change(dfn[x][d],1);
up[a[x]].change(dfo[x][d],-1);
sub[a[x]].change(dfn[x][d],1);
}
void build(int x,int d){
dfs1(x,x);
x=dfs2(x,x,sz[x]);
vis[x]=true;
cf[x][d]=x;
dep[x]=d;
dfn[x][d]=++ct;
for(int p=h[x];p;p=nxt[p]){
if(vis[v[p]]) continue;
dfs3(v[p],x,d,x,v[p]);
}
for(int p=h[x];p;p=nxt[p]){
if(vis[v[p]]) continue;
ans+=pa[v[p]][d][0]*1ll*(c[x][2]-cnt[v[p]][d][2]);
ans+=pa[v[p]][d][2]*1ll*(c[x][0]-cnt[v[p]][d][0]);
pd[x]+=(c[x][0]-cnt[v[p]][d][0])*1ll*cnt[v[p]][d][2];
pd[x]+=(c[x][2]-cnt[v[p]][d][2])*1ll*cnt[v[p]][d][0];
}
dfo[x][d]=++ct;
if(a[x]==0) ans+=sp[x][2];
if(a[x]==1) ans+=pd[x]/2;
if(a[x]==2) ans+=sp[x][0];
for(int p=h[x];p;p=nxt[p]) if(!vis[v[p]]) build(v[p],d+1);
}
void init(int n,vector<int> f,vector<int> u,vector<int> v,int q){
::n=n;
for(int i=0;i+1<n;i++) addedge(u[i],v[i]);
for(int i=0;i<n;i++) a[i]=f[i];
build(0,0);
}
void change(int x,int y){
if(a[x]==y) return;
for(int d=0;d<=dep[x];d++){
if(cf[x][d]==x){
if(a[x]==0) ans-=sp[x][2];
if(a[x]==1) ans-=pd[x]/2;
if(a[x]==2) ans-=sp[x][0];
if(y==0) ans+=sp[x][2];
if(y==1) ans+=pd[x]/2;
if(y==2) ans+=sp[x][0];
continue;
}
if(a[cf[x][d]]==0) ans-=sp[cf[x][d]][2];
if(a[cf[x][d]]==1) ans-=pd[cf[x][d]]/2;
if(a[cf[x][d]]==2) ans-=sp[cf[x][d]][0];
if(a[x]==1){
ans-=sub[0].ask(dfn[x][d]+1,dfo[x][d])*1ll*(c[cf[x][d]][2]-cnt[cr[x][d]][d][2]);
ans-=sub[2].ask(dfn[x][d]+1,dfo[x][d])*1ll*(c[cf[x][d]][0]-cnt[cr[x][d]][d][0]);
sp[cf[x][d]][0]-=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
sp[cf[x][d]][2]-=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
pa[cr[x][d]][d][0]-=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
pa[cr[x][d]][d][2]-=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
}else{
ans-=up[1].ask(dfn[x][d])*1ll*(c[cf[x][d]][2-a[x]]-cnt[cr[x][d]][d][2-a[x]]);
ans-=sp[cf[x][d]][2-a[x]]-pa[cr[x][d]][d][2-a[x]];
sp[cf[x][d]][a[x]]-=up[1].ask(dfn[x][d]);
pa[cr[x][d]][d][a[x]]-=up[1].ask(dfn[x][d]);
pd[cf[x][d]]-=(c[cf[x][d]][2-a[x]]-cnt[cr[x][d]][d][2-a[x]])*2;
}
c[cf[x][d]][a[x]]-=1;
cnt[cr[x][d]][d][a[x]]-=1;
sub[a[x]].change(dfn[x][d],-1);
up[a[x]].change(dfn[x][d],-1);
up[a[x]].change(dfo[x][d],+1);
if(y==1){
ans+=sub[0].ask(dfn[x][d]+1,dfo[x][d])*1ll*(c[cf[x][d]][2]-cnt[cr[x][d]][d][2]);
ans+=sub[2].ask(dfn[x][d]+1,dfo[x][d])*1ll*(c[cf[x][d]][0]-cnt[cr[x][d]][d][0]);
sp[cf[x][d]][0]+=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
sp[cf[x][d]][2]+=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
pa[cr[x][d]][d][0]+=sub[0].ask(dfn[x][d]+1,dfo[x][d]);
pa[cr[x][d]][d][2]+=sub[2].ask(dfn[x][d]+1,dfo[x][d]);
}else{
ans+=up[1].ask(dfn[x][d])*1ll*(c[cf[x][d]][2-y]-cnt[cr[x][d]][d][2-y]);
ans+=sp[cf[x][d]][2-y]-pa[cr[x][d]][d][2-y];
sp[cf[x][d]][y]+=up[1].ask(dfn[x][d]);
pa[cr[x][d]][d][y]+=up[1].ask(dfn[x][d]);
pd[cf[x][d]]+=(c[cf[x][d]][2-y]-cnt[cr[x][d]][d][2-y])*2;
}
c[cf[x][d]][y]+=1;
cnt[cr[x][d]][d][y]+=1;
sub[y].change(dfn[x][d],+1);
up[y].change(dfn[x][d],+1);
up[y].change(dfo[x][d],-1);
if(a[cf[x][d]]==0) ans+=sp[cf[x][d]][2];
if(a[cf[x][d]]==1) ans+=pd[cf[x][d]]/2;
if(a[cf[x][d]]==2) ans+=sp[cf[x][d]][0];
}
a[x]=y;
}
ll num_tours(){return ans;}
