#include "joitour.h"
#include <bits/stdc++.h>
using namespace std;
using i32 = int32_t;
#define int long long
const int N = 2e5 + 5;
const int LG = 20;
int n, bl[N], freq[N], sz[N], tot, timer = 0, tin[N], tout[N], f[N], ans[N], total = 0;
vector<int> adj[N], anc[N];
map<int, pair<int, int>> cd[N];
int par[N][20];
map<pair<int, int>, array<int, 5>> mp; // cent, child, 0, 1, 2, 01, 21
array<int, 8> info[N]; // s0, s1, s2, s01, s21, p(0,21), p(2,01), p(0,2);
struct Fenwick {
int maxn;
vector<vector<int>> fen;
void init(int sz) {
fen.resize(3);
maxn = sz;
for (int i = 0; i < 3; i++) fen[i].assign(maxn + 1, 0);
}
void upd(int f, int id, int k) {
for (; id <= maxn; id += (id & -id)) fen[f][id] += k;
}
int pref(int f, int id) {
int res = 0;
for (; id; id -= (id & -id)) res += fen[f][id];
return res;
}
int query(int f, int l, int r) {
return pref(f, r) - pref(f, l - 1);
}
} fen[N];
void get_sz(int u, int p) {
sz[u] = 1;
for (auto& v : adj[u]) {
if (v == p || bl[v]) continue;
get_sz(v, u);
sz[u] += sz[v];
}
}
int get_cent(int u, int p) {
int mx = 0, res = -1;
for (auto& v : adj[u]) {
if (v == p || bl[v]) continue;
res = max(res, get_cent(v, u));
mx = max(mx, sz[v]);
}
mx = max(mx, tot - sz[u]);
if (mx <= tot / 2) return u;
return res;
}
vector<int> node, tmp;
void dfs(int u, int p, int cent) {
tin[u] = ++timer;
// cerr << u << " pb " << cent << '\n';
anc[u].push_back(cent);
node.push_back(u);
for (auto& v : adj[u]) {
if (v == p || bl[v]) continue;
dfs(v, u, cent);
}
tout[u] = timer;
cd[cent][u] = {tin[u], tout[u]};
}
void dadfs(int u, int p, int rt, int cent) {
tmp.push_back(u);
par[u][anc[u].size() - 1] = rt;
for (auto& v : adj[u]) {
if (v == p || bl[v] || v == cent) continue;
dadfs(v, u, rt, cent);
}
}
void calc(int u) {
ans[u] = 0;
if (f[u] == 0) ans[u] += info[u][4];
if (f[u] == 1) ans[u] += info[u][0] * info[u][2] - info[u][7];
if (f[u] == 2) ans[u] += info[u][3];
ans[u] += info[u][0] * info[u][4] - info[u][5];
ans[u] += info[u][2] * info[u][3] - info[u][6];
}
int cnter = 0;
void dnc(int u) {
get_sz(u, -1);
tot = sz[u];
cnter += tot;
int cent = get_cent(u, -1);
// cerr << "cent: " << cent << '\n';
timer = 0;
node.clear();
dfs(cent, -1, cent);
fen[cent].init(timer);
for (auto& x : node) fen[cent].upd(f[x], tin[x], 1);
vector<vector<int>> ps(3, vector<int> (timer+1, 0));
for (auto& x : node) ps[f[x]][tin[x]]++;
for (int i = 0; i < 3; i++) for (int j = 1; j <= timer; j++) ps[i][j] += ps[i][j - 1];
for (int i = 0; i < 8; i++) info[cent][i] = 0;
for (auto& v : adj[cent]) if (!bl[v]) {
tmp.clear();
dadfs(v, -1, v, cent);
array<int, 5> arr = {0, 0, 0, 0, 0};
for (int i = 0; i < 3; i++) arr[i] = ps[i][tout[v]] - ps[i][tin[v] - 1];
for (auto& x : tmp) {
if (f[x] == 1) {
arr[3] += ps[0][tout[x]] - ps[0][tin[x] - 1];
arr[4] += ps[2][tout[x]] - ps[2][tin[x] - 1];
}
}
mp[{cent, v}] = arr;
// cerr << "made: " << cent << " " << v << '\n';
for (int i = 0; i < 5; i++) info[cent][i] += arr[i];
info[cent][5] += arr[0] * arr[4];
info[cent][6] += arr[2] * arr[3];
info[cent][7] += arr[0] * arr[2];
}
calc(cent);
bl[cent] = 1;
for (auto& v : adj[cent]) if (!bl[v]) dnc(v);
}
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]);
}
for (int i = 0; i < n; i++) f[i] = F[i];
dnc(0);
for (int i = 0; i < n; i++) total += ans[i];
}
void change(i32 X, i32 Y) {
int idx = 0;
for (auto& cent : anc[X]) {
total -= ans[cent];
int v = par[X][idx];
if (cent != X) {
// cout << "try: " << cent << " " << X << " " << v << '\n';
assert(mp.count({cent, v}));
array<int, 5> arr = mp[{cent, v}];
pair<int, int> tx = cd[cent][X];
info[cent][5] -= arr[0] * arr[4];
info[cent][6] -= arr[2] * arr[3];
info[cent][7] -= arr[0] * arr[2];
arr[f[X]] -= 1, info[cent][f[X]]--;
arr[Y] += 1, info[cent][Y]++;
info[cent][3] -= arr[3], info[cent][4] -= arr[4];
if (f[X] == 0 && Y == 1) {
arr[3] -= (arr[1] - 1) - (fen[cent].query(1, tx.first, tx.second));
arr[3] += fen[cent].query(0, tx.first, tx.second) - 1;
arr[4] += fen[cent].query(2, tx.first, tx.second);
} else if (f[X] == 1 && Y == 0) {
arr[3] -= fen[cent].query(0, tx.first, tx.second);
arr[4] -= fen[cent].query(2, tx.first, tx.second);
arr[3] += (arr[1] + 1) - fen[cent].query(1, tx.first, tx.second);
}
if (f[X] == 0 && Y == 2) {
arr[3] -= (arr[1]) - (fen[cent].query(1, tx.first, tx.second));
arr[4] += (arr[1]) - (fen[cent].query(1, tx.first, tx.second));
} else if (f[X] == 2 && Y == 0) {
arr[3] += (arr[1]) - (fen[cent].query(1, tx.first, tx.second));
arr[4] -= (arr[1]) - (fen[cent].query(1, tx.first, tx.second));
}
if (f[X] == 2 && Y == 1) {
arr[4] -= (arr[1] - 1) - (fen[cent].query(1, tx.first, tx.second));
arr[4] += fen[cent].query(2, tx.first, tx.second) - 1;
arr[3] += fen[cent].query(0, tx.first, tx.second);
} else if (f[X] == 1 && Y == 2) {
arr[3] -= fen[cent].query(0, tx.first, tx.second);
arr[4] -= fen[cent].query(2, tx.first, tx.second);
arr[4] += (arr[1] + 1) - fen[cent].query(1, tx.first, tx.second);
}
fen[cent].upd(f[X], tx.first, -1);
fen[cent].upd(Y, tx.second, 1);
info[cent][3] += arr[3], info[cent][4] += arr[4];
info[cent][5] += arr[0] * arr[4];
info[cent][6] += arr[2] * arr[3];
info[cent][7] += arr[0] * arr[2];
mp[{cent, v}] = arr;
// cout << cent << " " << v << " " << arr[0] << " " << arr[1] << " " << arr[2] << " " << arr[3] << " " << arr[4] << '\n';
} else {
f[X] = Y;
}
calc(cent);
// cout << cent << " " << ans[cent] << '\n';
total += ans[cent];
idx++;
}
f[X] = Y;
}
int num_tours() {
return total;
}
