# | Time | Username | Problem | Language | Result | Execution time | Memory |
---|---|---|---|---|---|---|---|
1091557 | thangdz2k7 | JOI tour (JOI24_joitour) | C++17 | 0 ms | 0 KiB |
This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
// author : thembululquaUwU
// 3.9.2024
#include <bits/stdc++.h>
#define pb push_back
#define fi first
#define se second
#define endl '\n'
using namespace std;
using ll = long long;
using ii = pair <int, int>;
using vi = vector <int>;
const int MaxN = 2e5 + 5;
const int mod = 1e9 + 7;
void maxl(auto &a, auto b) {a = max(a, b);}
void minl(auto &a, auto b) {a = min(a, b);}
int par[MaxN], sz[MaxN];
vi badj[MaxN], f;
int dfs(int u, int p = -1){
sz[u] = 1;
for (int v : badj[u]) if (v != p && par[v] < -1) sz[u] += dfs(v, u);
return sz[u];
}
int get_cen(int u, int m, int p = -1){
for (int v : badj[u]) if (v != p && par[v] < -1 && sz[v] * 2 > m) return get_cen(v, m, u);
return u;
}
int root;
void cen(int u = 0, int p = -1){
u = get_cen(u, dfs(u)); par[u] = p; if (p == -1) root = u;
for (int v : badj[u]) if (par[v] < -1) cen(v, u);
}
ll cnt[MaxN][8]; int sign[MaxN][8];
vi up = {1, 2, 4, 3, 6, 7};
vector <ii> pairs2 = {{2, 4}, {4, 2}, {2, 1}, {1, 2}};
vector <ii> pairs3 = {{1, 6}, {6, 1}, {3, 4}, {4, 3}};
void rsz(int u, int child){
if (u == -1) return;
rsz(par[u], u);
for (int t : up) cnt[u][t] -= cnt[child][t];
ll sub = cnt[u][1] * cnt[child][4] + cnt[u][4] * cnt[child][1];
cnt[u][5] -= sub; //if (f[u] == 2) cnt[u][7] -= sub;
for (auto [l, r] : pairs2) if (l == f[u]) cnt[u][l | r] -= cnt[child][r] * cnt[u][l];
for (auto [l, r] : pairs3) cnt[u][l | r] -= cnt[u][l] * cnt[child][r];
}
void rks(int u, int child){
if (u == -1) return;
for (auto [l, r] : pairs3) cnt[u][l | r] += cnt[u][l] * cnt[child][r];
for (auto [l, r] : pairs2) if (l == f[u]) cnt[u][l | r] += cnt[child][r] * cnt[u][l];
ll sub = cnt[u][1] * cnt[child][4] + cnt[u][4] * cnt[child][1];
cnt[u][5] += sub; //if (f[u] == 2) cnt[u][7] += sub;
for (int t : up) cnt[u][t] += cnt[child][t];
rks(par[u], u);
}
void add(int u, int t, int val){
rsz(par[u], u);
if (val < 0){
cnt[u][t] --;
ll sub = 0;
if (t == 1) sub = cnt[u][4];
if (t == 4) sub = cnt[u][1];
cnt[u][5] -= sub;
if (t == 2) cnt[u][7] -= cnt[u][5];
for (auto [l, r] : pairs2) if (l == f[u]) cnt[u][l | r] -= cnt[u][r];
for (auto [l, r] : pairs3) if (l == f[u]) cnt[u][l | r] -= cnt[u][r];
}
else {
for (auto [l, r] : pairs2) if (l == f[u]) cnt[u][l | r] += cnt[u][r];
for (auto [l, r] : pairs3) if (l == f[u]) cnt[u][l | r] += cnt[u][r];
ll sub = 0;
if (t == 1) sub = cnt[u][4];
if (t == 4) sub = cnt[u][1];
cnt[u][5] += sub;
if (t == 2) cnt[u][7] += cnt[u][5];
cnt[u][t] ++;
}
rks(par[u], u);
}
void init(int N, vi F, vi U, vi V, int Q){
for (int i = 0; i < N - 1; ++ i){
int u = U[i], v = V[i];
badj[u].pb(v); badj[v].pb(u);
}
for (int i = 0; i < N; ++ i) par[i] = -2; cen();
f.resize(N, 0);
for (int i = 0; i < N; ++ i) {
f[i] = (1 << F[i]);
add(i, f[i], 1);
}
//f[0] = (1 << F[0]); add(0, f[0], 1);
//f[2] = (1 << F[2]); add(2, f[2], 1);
//f[1] = (1 << F[1]); add(1, f[1], 1);
//for (int i = 0; i < N; ++ i){
//for (int j = 0; j < 8; ++ j) cout << cnt[i][j] << ' ';
//cout << endl;
//}
}
void change(int u, int t){
add(u, f[u], -1); t = 1 << t;
f[u] = t; add(u, f[u], 1);
}
ll num_tours(){ return cnt[root][7];}
void solve(){
int n; cin >> n; vi f(n); for (int i = 0; i < n; ++ i) cin >> f[i];
vi u(n - 1), v(n - 1); for (int i = 0; i < n - 1; ++ i) cin >> u[i] >> v[i];
int q; cin >> q; init(n, f, u, v, q);
cout << num_tours() << endl;
for (int i = 1; i <= q; ++ i){
int x, y; cin >> x >> y;
change(x, y); cout << num_tours() << endl;
}
}
int main(){
if (fopen("pqh.inp", "r")){
freopen("pqh.inp", "r", stdin);
freopen("pqh.out", "w", stdout);
}
ios_base::sync_with_stdio(0);
cin.tie(0); cout.tie(0);
int t = 1; // cin >> t;
while (t --) solve();
return 0;
}