#include "joitour.h"
#include <bits/stdc++.h>
using ll = long long ;
using ull = unsigned ll;
#define f first
#define s second
#define pb push_back
#define epb emplace_back
using namespace std;
using pii = pair<int,int>;
using pll = pair<ll,ll>;
struct Data{
vector <ll> f[3];
int n;
vector <ll> cnt;
vector <vector<ll>>cnts;
ll s_bad;
Data() {
cnt.resize(5);
}
void add(int pos, int val, int id) {
while(pos < f[id].size()) {
f[id][pos] += val;
pos += (pos & -pos);
}
}
ll get(int pos, int id) {
ll res = 0;
while(pos) {
res += f[id][pos];
pos -= (pos & -pos);
}
return res;
}
ll get(int l, int r, int id) {
return get(r, id) - get(l - 1, id);
}
};
const int nmax = 200001;
vector <vector <int>>vec(nmax);
vector <vector <int>>sz(nmax);
vector <vector<int>>in(nmax);
vector <vector <int>>g(nmax);
vector <int>y(nmax);
vector <vector<int>>lead_in(nmax);
Data dt[nmax];
vector <bool> blocked(nmax);
vector <int> siz(nmax);
ll ans;
int timer = 1;
vector <vector<ll>>cnt(nmax, vector<ll>(5));
void dfs(int v, int e) {
// cout << v << ' ';
siz[v] = 1;
for(int to : g[v]) {
if(to == e || blocked[to]) continue;
dfs(to, v);
siz[v] += siz[to];
}
}
int find_cent(int v, int e, int n) {
for(int to : g[v]) {
if(blocked[to] || to == e) continue;
if(siz[to] * 2 >= n) return find_cent(to, v, n);
}
return v;
}
void clear_dfs(int v, int e) {
siz[v] = 0;
for(int to : g[v]) {
if(to == e || blocked[to]) continue;
clear_dfs(to, v);
}
siz[v] = 0;
}
void update_cnt(int v, int to, bool x = false) {
for(int i = 0; i < 5; i++) cnt[v][i] += cnt[to][i];
if(y[v] == 1 && !x) {
cnt[v][3] += cnt[to][0];
cnt[v][4] += cnt[to][2];
}
}
void DFS(int v, int e, int num, vector<int>&visited) {
in[v].pb(timer++);
lead_in[v].pb(num);
sz[v].pb(1);
visited.pb(v);
for(int j = 0; j < 5; j++) cnt[v][j] = 0;
for(int to : g[v]) {
if(to == e || blocked[to]) continue;
DFS(to, v, num, visited);
update_cnt(v, to);
sz[v].back() += sz[to].back();
}
cnt[v][y[v]]++;
}
void clearDFS(int v, int e) {
cnt[v].assign(5, 0);
for(int to : g[v]) {
if(to == e || blocked[to]) continue;
clearDFS(to, v);
}
}
void decompose(int v) {
dfs(v, v);
int n = siz[v];
int x = find_cent(v, v, siz[v]);
clear_dfs(v, v);
timer = 1;
sz[x].pb(n);
in[x].pb(1);
vec[x].pb(x);
lead_in[x].pb(0);
int num = 0;
vector <int> visited;
for(int to : g[x]) {
if(blocked[to]) continue;
DFS(to, x, num, visited);
update_cnt(x, to, true);
dt[x].cnts.pb(cnt[to]);
dt[x].s_bad += cnt[to][2] * cnt[to][0];
num++;
}
cnt[x][y[x]]++;
for(int j = 0; j< 3; j++) {
dt[x].f[j].resize(n + 1);
}
dt[x].cnt = cnt[x];
for(auto to : visited) {
vec[to].pb(x);
if(y[to] != 1) {
dt[x].add(in[to].back(), 1, y[to]);
}
else {
dt[x].add(in[to].back(), 1, 1);
dt[x].add(in[to].back() + sz[to].back(), -1, y[to]);
}
}
clearDFS(x, x);
timer = 1;
blocked[x] = true;
for(int to : g[x]) {
if(!blocked[to]) decompose(to);
}
}
/*
ans += (dt[i].cnt[0] - to[0]) * to[4] +
(dt[i].cnt[2] - to[2]) * to[3] +
to[0] * (dt[i].cnt[4] - to[4]) +
to[2] * (dt[i].cnt[3] - to[3]);
if(y[i] == 1) {
ans += to[2] * (dt[i].cnt[0] - to[0]) + to[0] * (dt[i].cnt[2] - to[2]);
}
if(y[i] == 0) {
ans += to[4];
}
if(y[i] == 2) {
ans += to[3];
}
*/
void updCent0(int v, int id, ll sign) {
ans += 2 * sign * (dt[v].cnt[4]);
dt[v].cnt[0] += sign;
}
void update0(int v, int id, ll sign) {
int x = vec[v][id];
int l = lead_in[v][id];
vector <ll>&to = dt[x].cnts[l];
if(v == x) {
updCent0(v, id, sign);
return;
}
int intm = in[v][id];
int sztm = sz[v][id];
ll cnt = dt[x].get(intm, 1);
ans += 2 * sign * (dt[x].cnt[2] - to[2]) * cnt;
ans += 2 * sign * (dt[x].cnt[4] - to[4]);
if(y[x] == 1) {
ans += 2 * sign * (dt[x].cnt[2] - to[2]);
}
to[0] += sign; dt[x].cnt[0] += sign;
to[3] += sign * cnt; dt[x].cnt[3] += cnt * sign;
dt[x].s_bad += sign * to[2];
dt[x].add(intm, sign, 0);
}
void updCent2(int v, int id, ll sign) {
ans += 2 * sign * dt[v].cnt[3];
dt[v].cnt[2] += sign;
}
void update2(int v, int id, ll sign) {
int x = vec[v][id];
if(v == x) {
updCent2(v, id, sign);
return;
}
int l = lead_in[v][id];
vector <ll>&to = dt[x].cnts[l];
int intm = in[v][id];
int sztm = sz[v][id];
ll cnt = dt[x].get(intm, 1);
ans += 2 * sign * (dt[x].cnt[0] - to[0]) * cnt;
ans += 2 * sign * (dt[x].cnt[3] - to[3]);
if(y[x] == 1) {
ans += 2 * sign * (dt[x].cnt[0] - to[0]);
}
to[2] += sign; dt[x].cnt[2] += sign;
to[4] += sign * cnt; dt[x].cnt[4] += sign * cnt;
dt[x].s_bad += sign * to[0];
dt[x].add(intm, sign, 2);
}
void updCent1(int v, ll sign) {
ans += 2 * sign * (dt[v].cnt[0] * dt[v].cnt[2] - dt[v].s_bad);
dt[v].cnt[1] += sign;
}
void update1(int v, int id, ll sign) {
int x = vec[v][id];
if(v == x) {
updCent1(v, sign);
return;
}
int l = lead_in[v][id];
vector <ll>&to = dt[x].cnts[l];
int intm = in[v][id];
int sztm = sz[v][id];
//cout << x << ' ' << intm << ' ' << sztm << ' ';
ll num0 = dt[x].get(intm, intm + sztm - 1, 0);
// cout << num0 << ' ' ;
ll num2 = dt[x].get(intm, intm + sztm - 1, 2);
// cout << num2 << '\n';
ans += sign * 2 * num0 * (dt[x].cnt[2] - to[2]);
ans += sign * 2 * num2 * (dt[x].cnt[0] - to[0]);
to[1] += sign; to[3] += num0 * sign; to[4] += num2 * sign;
dt[x].cnt[1] += sign; dt[x].cnt[3] += num0 * sign; dt[x].cnt[4] += num2 * sign;
dt[x].add(intm, sign, 1);
dt[x].add(intm + sztm, -sign, 1);
}
void init(int n, std::vector<int> f, std::vector<int> u, std::vector<int> v,int q) {
for(int i = 0; i < f.size(); i++) {
y[i] = f[i];
}
// cout << 1;
for(int i = 0; i < u.size(); i++) {
g[u[i]].pb(v[i]);
g[v[i]].pb(u[i]);
}
// cout << 1;
decompose(0);
// cout << 1;
for(int i = 0; i < n; i++) {
for(auto to : dt[i].cnts) {
ans += (dt[i].cnt[0] - to[0]) * to[4] + (dt[i].cnt[2] - to[2]) * to[3] + to[0] * (dt[i].cnt[4] - to[4]) + to[2] * (dt[i].cnt[3] - to[3]);
if(y[i] == 1) {
ans += to[2] * (dt[i].cnt[0] - to[0]) + to[0] * (dt[i].cnt[2] - to[2]);
}
if(y[i] == 0) {
ans += to[4];
}
if(y[i] == 2) {
ans += to[3];
}
}
// cout << ans << ' ';
}
}
void change(int X, int Y) {
for(int i = 0; i < vec[X].size(); i++) {
if(y[X] == 0) {
update0(X, i, -1);
}
if(y[X] == 1)
update1(X, i, -1);
if(y[X] == 2)
update2(X, i, -1);
}
// return;
y[X] = Y;
for(int i = 0; i < vec[X].size(); i++) {
if(y[X] == 0) {
update0(X, i, 1);
}
if(y[X] == 1)
update1(X, i, 1);
if(y[X] == 2)
update2(X, i, 1);
}
}
long long num_tours() {
return ans / 2;
}
