#include "joitour.h"
#include<algorithm>
#include<vector>
#include<utility>
#include<algorithm>
using namespace std;
typedef long long ll;
namespace{
int n;
ll cnta = 0, cntb = 0, cntc = 0, allp;
vector<int> cur, head, in, out, ord, sz, par, wid;
vector<pair<int, int>> hans, uans;
int timer = 0;
vector<vector<int>> graph;
vector<vector<pair<int, int>>> wans;
void dfs(int node, int parent){
sz[node] = 1;
par[node] = parent;
head[node] = node;
int maxi = 0, pos = -1;
for(auto &x: graph[node]){
if(x == parent) continue;
dfs(x, node);
sz[node] += sz[x];
if(sz[x] > maxi){
maxi = sz[x];
pos = x;
}
}
if(pos > 0) head[pos] = node;
for(auto &x: graph[node]){
if(x == pos) swap(x, graph[node][0]);
}
}
void dfs_head(int node, int parent){
ord[++timer] = node;
in[node] = timer;
if(head[node] != node) head[node] = head[parent];
int i = -1;
for(auto &x: graph[node]){
i++;
if(x == parent) continue;
wid[x] = i;
dfs_head(x, node);
uans[node].first += uans[x].first;
uans[node].second += uans[x].second;
}
if(cur[node] == 0) uans[node].first++;
else if(cur[node] == 2) uans[node].second++;
i = -1;
for(auto &x: graph[node]){
i++;
if(x == parent) continue;
if(i == 0) hans[node] = uans[x];
else wans[node][i] = uans[x];
}
out[node] = timer;
}
struct node{
ll ca, cc, acta, actc, act, sum;
node(ll _ca, ll _cc, ll _acta, ll _actc, ll _act, ll _sum){
ca = _ca, cc = _cc, acta = _acta, actc = _actc, sum = _sum, act = _act;
}
};
struct st{
vector<node> seg;
};
}
void init(int N, std::vector<int> F, std::vector<int> U, std::vector<int> V, int Q) {
n = N;
cur.resize(n + 1);
graph.resize(n + 1);
head.resize(n + 1);
in.resize(n + 1);
out.resize(n + 1);
ord.resize(n + 1);
sz.resize(n + 1);
par.resize(n + 1);
hans.resize(n + 1);
uans.resize(n + 1);
wid.resize(n + 1);
wans.resize(n + 1);
for(int i = 0; i < n; i++){
cur[i + 1] = F[i];
if(F[i] == 0) cnta++;
else if(F[i] == 1) cntb++;
else cntc++;
}
allp = cnta * cntb * cntc;
for(int i = 0; i < n - 1; i++){
U[i]++;
V[i]++;
graph[U[i]].push_back(V[i]);
graph[V[i]].push_back(U[i]);
}
for(int i = 1; i <= n; i++){
wans[i].resize((int)graph[i].size());
}
dfs(1, 1);
dfs_head(1, 1);
for(int i = 1; i <= n; i++){
uans[i].first = cnta - uans[i].first;
uans[i].second = cntc - uans[i].second;
}
}
void change(int X, int Y) {
}
long long num_tours() {
ll ans = 0;
for(int i = 1; i <= n; i++){
ans += 1ll * uans[i].first * uans[i].second;
ans += 1ll * hans[i].first * hans[i].second;
int sz = (int)graph[i].size();
for(int j = 0; j < sz; j++) ans += 1ll * wans[i][j].first * wans[i][j].second;
}
return allp - ans;
}
// g++ -std=c++17 -Wall -Wextra -Wshadow -fsanitize=undefined -fsanitize=address -o run grader.cpp joitour.cpp
