#include "joitour.h"
#include <bits/stdc++.h>
using namespace std;
using i32 = int32_t;
#define int long long
const int maxn = 2e5 + 5;
vector<int> adj[maxn];
int n, f[maxn], cnt[maxn][3], par[maxn], c0, c2, sc0, sc2, sum;
void dfs(int u, int p = -1) {
for (int i = 0; i < 3; i++) cnt[u][i] = 0;
cnt[u][f[u]] = 1;
for (auto& v : adj[u]) {
if (v == p) continue;
par[v] = u;
dfs(v, u);
for (int i = 0; i < 3; i++) {
cnt[u][i] += cnt[v][i];
}
}
}
void init(i32 N, std::vector<i32> F, std::vector<i32> U, std::vector<i32> V,
i32 Q) {
n = N;
for (int i = 0; i < N - 1; i++) {
adj[U[i]].push_back(V[i]);
adj[V[i]].push_back(U[i]);
}
par[0] = -1;
for (int i = 0; i < N; i++) f[i] = F[i], c0 += (f[i] == 0), c2 += (f[i] == 2);
dfs(0, -1);
for (int i = 0; i < n; i++) {
if (f[i] == 1) {
int u = i;
sc0 += cnt[u][0];
sc2 += cnt[u][2];
sum -= cnt[u][0] * cnt[u][2];
for (auto& v : adj[u]) if (v != par[u]) sum -= cnt[v][0] * cnt[v][2];
}
}
}
void change(i32 X, i32 Y) {
c0 -= (f[X] == 0);
c2 -= (f[X] == 2);
int u = X, x = u;
while (par[x] != -1) {
x = par[x];
if (f[x] == 1) {
sc0 -= cnt[x][0];
sc2 -= cnt[x][2];
sum += cnt[x][0] * cnt[x][2];
for (auto& v : adj[x]) if (v != par[x]) sum += cnt[v][0] * cnt[v][2];
}
}
if (f[X] == 1) {
sc0 -= cnt[u][0];
sc2 -= cnt[u][2];
sum += cnt[u][0] * cnt[u][2];
for (auto& v : adj[u]) if (v != par[u]) sum += cnt[v][0] * cnt[v][2];
}
// cout << "SUM " << sum << '\n';
cnt[u][f[X]] -= 1;
cnt[u][Y] += 1;
if (Y == 1) {
sc0 += cnt[u][0];
sc2 += cnt[u][2];
sum -= cnt[u][0] * cnt[u][2];
for (auto& v : adj[u]) if (v != par[u]) sum -= cnt[v][0] * cnt[v][2];
}
while (par[u] != -1) {
u = par[u];
cnt[u][f[X]] -= 1;
cnt[u][Y] += 1;
if (f[u] == 1) {
sc0 += cnt[u][0];
sc2 += cnt[u][2];
sum -= cnt[u][0] * cnt[u][2];
for (auto& v : adj[u]) if (v != par[u]) sum -= cnt[v][0] * cnt[v][2];
}
}
f[X] = Y;
c0 += (f[X] == 0);
c2 += (f[X] == 2);
}
int num_tours() {
int ans = 0;
// cout << c0 << " " << sc2 << " " << c2 << " " << sc0 << " " << sum << '\n';
return c0 * sc2 + c2 * sc0 + sum;
}
