#include "joitour.h"
#include <bits/stdc++.h>
using namespace std;
vector<vector<int>> adj;
vector<pair<long long,long long>> subtot;
vector<int> type;
vector<int> p;
vector<long long> tot;
long long ans;
long long sigma0;
long long sigma2;
void process(int x,long long offset){
if(type[x]!=1){
pair<long long,long long> mytot = {0,0};
if(type[x]==0)mytot.first++;
else mytot.second++;
for(int&i:adj[x])if(i!=p[x]){
auto t = subtot[i];
mytot.first+=t.first;
mytot.second+=t.second;
}
subtot[x]=mytot;
return;
}
pair<long long,long long> mytot = {0,0};
for(int&i:adj[x])if(i!=p[x]){
auto t = subtot[i];
ans-=t.first*t.second*offset;
mytot.first+=t.first;
mytot.second+=t.second;
}
ans-=mytot.first*mytot.second*offset;
sigma0+=mytot.first*offset;
sigma2+=mytot.second*offset;
subtot[x]=mytot;
};
void init(int N, vector<int> F, vector<int> U, vector<int> V,
int Q) {
adj = vector<vector<int>> (N);
type = F;
p = vector<int>(N);
tot = vector<long long>(3);
subtot = vector<pair<long long,long long>>(N);
for(int i=0;i<N-1;i++){
adj[U[i]].emplace_back(V[i]);
adj[V[i]].emplace_back(U[i]);
}
for(int&i:type)tot[i]++;
ans = sigma0 = sigma2 = 0;
function<void(int,int)> dfs = [&](int x,int par){
p[x]=par;
for(int&i:adj[x])if(i!=par)dfs(i,x);
};
dfs(0,-1);
for(int i=N-1;i>=0;i--)process(i,1);
}
void change(int X, int Y){
int backX = X;
while(X!=-1){
process(X,-1);
X = p[X];
}
X = backX;
tot[type[X]]--;
type[X]=Y;
tot[type[X]]++;
while(X!=-1){
process(X,1);
X = p[X];
}
}
long long num_tours(){
return tot[2]*sigma0+tot[0]*sigma2+ans;
}
