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