#include "joitour.h"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int N=2e5+5;
int n;
vector<int> adj[N];
int a[N];
enum Type{Compress,Rake,AddEdge,AddVertex,Vertex};
struct StaticTopTree{
int root,node_id;
int hv[N],sz[N],p[N];
int lch[4*N],rch[4*N],par[4*N];
Type type[4*N];
void dfs(int u){
sz[u]=1;
for(auto v:adj[u])if(v!=p[u]){
p[v]=u;
dfs(v);
sz[u]+=sz[v];
if(sz[v]>sz[hv[u]])hv[u]=v;
}
}
int add(int i,int l,int r,Type t){
if(!i)i=++node_id;
lch[i]=l,rch[i]=r,type[i]=t;
if(l)par[l]=i;
if(r)par[r]=i;
return i;
}
pair<int,int> merge(const vector<pair<int,int>> &a,Type t){
if(a.size()==1)return a[0];
int tot=0;
vector<pair<int,int>> b,c;
for(auto [i,s]:a)tot+=s;
for(auto [i,s]:a){
(tot>s?b:c).emplace_back(i,s);
tot-=s*2;
}
auto [i,si]=merge(b,t);
auto [j,sj]=merge(c,t);
return {add(0,i,j,t),si+sj};
}
pair<int,int> compress(int i){
vector<pair<int,int>> a{add_vertex(i)};
while(hv[i])a.emplace_back(add_vertex(i=hv[i]));
return merge(a,Compress);
}
pair<int,int> rake(int i){
vector<pair<int,int>> a;
for(auto j:adj[i])if(j!=p[i]&&j!=hv[i])a.emplace_back(add_edge(j));
return a.empty()?make_pair(0,0):merge(a,Rake);
}
pair<int,int> add_edge(int i){
auto [j,s]=compress(i);
return {add(0,j,0,AddEdge),s};
}
pair<int,int> add_vertex(int i){
auto [j,s]=rake(i);
return {add(i,j,0,j?AddVertex:Vertex),s+1};
}
void build(){
node_id=n;
dfs(1);
root=compress(1).first;
}
}stt;
struct Point{
ll c0,c1,c2,c10,c12,c02;
ll ans;
}point[4*N];
struct Path{
ll c0,c1,c2,c10,c12;
ll p0,p1,p2,p10,p12;
ll s0,s1,s2;
ll ans;
}path[4*N];
Path compress(Path p,Path c){
Path res;
res.c0=p.c0+c.c0,res.c1=p.c1+c.c1,res.c2=p.c2+c.c2;
res.p0=p.p0+c.p0,res.p1=p.p1+c.p1,res.p2=p.p2+c.p2;
res.s0=p.s0+c.s0,res.s1=p.s1+c.s1,res.s2=p.s2+c.s2;
res.c10=p.c10+c.c10+p.s1*c.c0;
res.c12=p.c12+c.c12+p.s1*c.c2;
res.p10=p.p10+c.p10+c.s1*p.p0;
res.p12=p.p12+c.p12+c.s1*p.p2;
res.ans=p.ans+c.ans+p.p0*c.c12+p.p2*c.c10+c.c0*p.p12+c.c2*p.p10;
return res;
}
Point rake(Point l,Point r){
Point res;
res.c0=l.c0+r.c0;
res.c1=l.c1+r.c1;
res.c2=l.c2+r.c2;
res.c02=l.c0*r.c2+l.c2*r.c0;
res.c10=l.c10+r.c10;
res.c12=l.c12+r.c12;
res.ans=l.ans+r.ans+l.c0*r.c12+l.c2*r.c10+r.c0*l.c12+r.c2*l.c10;
return res;
}
Point add_edge(Path x){
Point res;
res.c0=x.c0,res.c1=x.c1,res.c2=x.c2;
res.c10=x.c10,res.c12=x.c12;
res.c02=0;
res.ans=x.ans;
return res;
}
Path add_vertex(Point x,int i){
Path res;
res.s0=a[i]==0,res.s1=a[i]==1,res.s2=a[i]==2;
res.c0=res.p0=x.c0+res.s0;
res.c1=res.p1=x.c1+res.s1;
res.c2=res.p2=x.c2+res.s2;
res.c10=res.p10=x.c10+res.s1*x.c0;
res.c12=res.p12=x.c12+res.s1*x.c2;
res.ans=x.ans+res.s1*x.c02+res.s0*x.c12+res.s2*x.c10;
return res;
}
Path vertex(int i){
Path res;
res.c0=res.p0=res.s0=a[i]==0;
res.c1=res.p1=res.s1=a[i]==1;
res.c2=res.p2=res.s2=a[i]==2;
res.c10=res.c12=res.p10=res.p12=res.ans=0;
return res;
}
void update(int i){
if(stt.type[i]==Compress)path[i]=compress(path[stt.lch[i]],path[stt.rch[i]]);
else if(stt.type[i]==Rake)point[i]=rake(point[stt.lch[i]],point[stt.rch[i]]);
else if(stt.type[i]==AddEdge)point[i]=add_edge(path[stt.lch[i]]);
else if(stt.type[i]==AddVertex)path[i]=add_vertex(point[stt.lch[i]],i);
else if(stt.type[i]==Vertex)path[i]=vertex(i);
}
void dfs(int i){
if(!i)return;
dfs(stt.lch[i]);
dfs(stt.rch[i]);
update(i);
}
void init(int _n,vector<int> _a,vector<int> _u,vector<int> _v,int Q){
n=_n;
for(int i=1;i<=n;i++)a[i]=_a[i-1];
for(int i=0;i<n-1;i++){
int u=_u[i]+1,v=_v[i]+1;
adj[u].emplace_back(v);
adj[v].emplace_back(u);
}
stt.build();
dfs(stt.root);
}
void change(int i,int t){
i++;
a[i]=t;
for(int u=i;u;u=stt.par[u])update(u);
}
ll num_tours(){
return path[stt.root].ans;
}
