#include "bits/stdc++.h" /** keep-include */
using namespace std;
#define rep(i, b, e) for(int i = (b); i <= (e); i++)
#define per(i, b, e) for(int i = (e); i >= (b); i--)
#define FOR(i, b, e) rep(i, b, (e) - 1)
#define SZ(x) int(x.size())
#define all(x) x.begin(), x.end()
#define pb push_back
#define mp make_pair
#define st first
#define nd second
using ll = long long;
using vi = vector<int>;
using pii = pair<int, int>;
auto &operator<<(auto &o, pair<auto, auto> p) {
return o << "(" << p.st << ", " << p.nd << ")"; }
auto operator<<(auto &o, auto x)->decltype(end(x), o) {
o << "{"; int i=0; for(auto e: x) o << ", " + 2*!i++ << e;
return o << "}"; }
#ifdef LOCAL
#define deb(x...) cerr << "[" #x "]: ", [](auto...$) { \
((cerr << $ << "; "),...) << endl; }(x)
#else
#define deb(...)
#endif
#include "joitour.h"
const int N=2e5+5;
ll result;
int typ[N];
struct drzewo{
vi real_name;
vi siz, par, pre, post, subtree, order;
vector<vi> E;
int root=-1, real_root;
int n;
vi ile[2];
vector<ll> sum[2];
ll iles[2]={}, sums[2]={};
ll zle=0;
ll res=0;
int new_name(int v){
return lower_bound(all(real_name), v)-real_name.begin();
}
void dfs_cent(int v){
siz[v]=1;
bool ok=1;
for(int u:E[v]){
if(u!=par[v]){
par[u]=v;
dfs_cent(u);
siz[v]+=siz[u];
if(siz[u]*2>n)ok=0;
}
}
if(ok && 2*(n-siz[v])<=n){
root=v;
}
}
void dfs_init(int v, int sub=-1){
deb(v, sub, order, root, E[v]);
pre[v]=order.size();
order.pb(v);
int cnt=0;
if(v!=root){
subtree[v]=sub;
for(int &u:E[v])if(u==par[v])swap(u, E[v].back());
E[v].pop_back();
}
for(int u:E[v]){
par[u]=v;
if(v==root)dfs_init(u, cnt);
else dfs_init(u, sub);
cnt++;
}
post[v]=order.size();
}
pair<int, ll> licz(int v, int lbl, int num){//num - liczba 1
pair<int, ll> a;
for(int u:E[v]){
auto b=licz(u, lbl, num+(typ[real_name[v]]==1));
a.st+=b.st;
a.nd+=b.nd;
}
if(typ[real_name[v]]==lbl*2){
a.st++;
a.nd+=num;
}
return a;
}
void calc(){
result-=res;
res=0;
iles[0]=0;
iles[1]=0;
sums[0]=0;
sums[1]=0;
zle=0;
for(int i=0; i<E[root].size(); i++){
auto a=licz(E[root][i], 0, 0);
auto b=licz(E[root][i], 1, 0);
ile[0][i]={a.st};
ile[1][i]={b.st};
sum[0][i]={a.nd};
sum[1][i]={b.nd};
iles[0]+=a.st;
iles[1]+=b.st;
sums[0]+=a.nd;
sums[1]+=b.nd;
res-=a.st*b.nd;
res-=a.nd*b.st;
zle+=a.st*1ll*b.st;
}
res+=sums[0]*iles[1];
res+=sums[1]*iles[0];
if(typ[real_name[root]]==1)res+=iles[0]*1ll*iles[1]-zle;
else res+=sums[1-typ[real_name[root]]/2];
result+=res;
}
vector<int> suma[2];
vector<int> srod;
void add_srod(int l, int r, int c){
for(l+=n, r+=n; l<r; l>>=1, r>>=1){
if(l&1)srod[l++]+=c;
if(r&1)srod[--r]+=c;
}
}
int get_srod(int v){
v+=n;
int ans=0;
while(v>0){
ans+=srod[v];
v/=2;
}
return ans;
}
int licz_suma(int l, int r, int ktora){
int ans=0;
for(l+=n, r+=n; l<r; l>>=1, r>>=1){
if(l&1)ans+=suma[ktora][l++];
if(r&1)ans+=suma[ktora][--r];
}
return ans;
}
void add_suma(int v, int c, int ktora){
v+=n;
while(v>0){
suma[ktora][v]+=c;
v/=2;
}
}
drzewo(vi V, vector<pii> edges){
deb(V, edges);
real_name=V;
sort(all(real_name));
n=real_name.size();
E.resize(n);
siz.resize(n);
par.resize(n);
subtree.resize(n);
pre.resize(n);
post.resize(n);
par[0]=-1;
for(auto &[u, v] : edges){
deb(u, v, real_name);
u=new_name(u);
v=new_name(v);
E[u].pb(v);
E[v].pb(u);
}
dfs_cent(0);
assert(root!=-1);
deb(root, real_name);
real_root=real_name[root];
par[root]=-1;
ile[0].resize(E[root].size());
ile[1].resize(E[root].size());
sum[0].resize(E[root].size());
sum[1].resize(E[root].size());
dfs_init(root);
calc();
suma[0].resize(n+n);
suma[1].resize(n+n);
srod.resize(n+n);
for(int i=0; i<n; i++){
if(i==root)continue;
if(typ[real_name[i]]==1){
add_srod(pre[i], post[i], 1);
}
else{
add_suma(pre[i], 1, typ[real_name[i]]/2);
}
}
}
pair<vi, vector<pii> > prep(int v){
vi vert;
vector<pii> edg;
for(int i=pre[v]; i<post[v]; i++){
int u=order[i];
vert.pb(real_name[u]);
for(int w:E[u]){
edg.pb(mp(real_name[u], real_name[w]));
}
}
return {vert, edg};
}
void update_slow(){
calc();
deb(res);
}
int licz_gora(int v){
deb(v, typ[real_name[v]]);
if(v==root)return 0;
return licz_gora(par[v])+(typ[real_name[v]]==1);
}
int licz_gora2(int v){
return get_srod(pre[v]);
}
int licz_dol(int v, int lbl){
int cnt=0;
for(int i=pre[v]; i<post[v]; i++){
if(typ[real_name[order[i]]]==lbl)cnt++;
}
return cnt;
}
int licz_dol2(int v, int lbl){
return licz_suma(pre[v], post[v], lbl/2);
}
void update_fast(int v, int stary, int nowy){
v=new_name(v);
deb(root, v, stary, nowy);
deb(ile[0], ile[1], sum[0], sum[1]);
result-=res;
typ[real_name[v]]=stary;
if(typ[real_name[root]]==1)res-=iles[0]*1ll*iles[1]-zle;
else res-=sums[1-typ[real_name[root]]/2];
if(v!=root){
res-=sums[0]*iles[1];
res-=sums[1]*iles[0];
int sub=subtree[v];
iles[0]-=ile[0][sub];
iles[1]-=ile[1][sub];
sums[0]-=sum[0][sub];
sums[1]-=sum[1][sub];
res+=ile[0][sub]*sum[1][sub];
res+=ile[1][sub]*sum[0][sub];
zle-=ile[1][sub]*1ll*ile[0][sub];
if(stary==1){
int cnt0=licz_dol2(v, 0);
int cnt2=licz_dol2(v, 2);
sum[0][sub]-=cnt0;
sum[1][sub]-=cnt2;
add_srod(pre[v], post[v], -1);
}
else{
int cnt=licz_gora2(v);
//deb(cnt, licz_gora2(v));
//assert(cnt==licz_gora2(v));
sum[stary/2][sub]-=cnt;
ile[stary/2][sub]--;
add_suma(pre[v], -1, typ[real_name[v]]/2);
}
deb(ile[0], ile[1], sum[0], sum[1]);
deb(res);
typ[real_name[v]]=nowy;
if(nowy==1){
int cnt0=licz_dol2(v, 0);
int cnt2=licz_dol2(v, 2);
sum[0][sub]+=cnt0;
sum[1][sub]+=cnt2;
add_srod(pre[v], post[v], 1);
}
else{
int cnt=licz_gora2(v);
//deb(cnt, licz_gora2(v));
//assert(cnt==licz_gora2(v));
sum[nowy/2][sub]+=cnt;
ile[nowy/2][sub]++;
add_suma(pre[v], 1, typ[real_name[v]]/2);
}
iles[0]+=ile[0][sub];
iles[1]+=ile[1][sub];
sums[0]+=sum[0][sub];
sums[1]+=sum[1][sub];
res-=ile[0][sub]*sum[1][sub];
res-=ile[1][sub]*sum[0][sub];
zle+=ile[1][sub]*1ll*ile[0][sub];
res+=sums[0]*iles[1];
res+=sums[1]*iles[0];
}
typ[real_name[v]]=nowy;
if(typ[real_name[root]]==1)res+=iles[0]*1ll*iles[1]-zle;
else res+=sums[1-typ[real_name[root]]/2];
result+=res;
deb(root, v, stary, nowy);
deb(ile[0], ile[1], sum[0], sum[1]);
deb(res);
}
};
vector<drzewo*> drzewa(N);
int rodzic[N];
void stworz(vector<int> vert, vector<pii> edg, int par){
drzewo* T=new drzewo(vert, edg);
drzewa[T->real_root]=T;
deb(T->real_root);
rodzic[T->real_root]=par;
for(int u:T->E[T->root]){
auto a = T->prep(u);
deb(T->real_root, a.st, a.nd);
stworz(a.st, a.nd, T->real_root);
}
}
int n;
void init(int _n, std::vector<int> F, std::vector<int> U, std::vector<int> V,
int Q) {
n=_n;
vector<pii> edges;
for(int i=0; i<n-1; i++){
edges.pb(mp(U[i], V[i]));
}
for(int i=0; i<n; i++){
typ[i]=F[i];
}
vi v(n);
iota(all(v), 0);
stworz(v, edges, -1);
}
void change(int X, int Y) {
int stary=typ[X];
int v=X;
typ[X]=Y;
while(X!=-1){
drzewa[X]->update_fast(v, stary, Y);
//drzewa[X]->update_slow();
X=rodzic[X];
}
}
long long num_tours() {
return result;
}
