#include "catdog.h"
#include<bits/stdc++.h>
using namespace std;
const int MX=1e5+5;
int n, sz[MX], par[MX], num[MX];
int hldpar[MX], pos[MX], st[MX];
vector<int>adj[MX];
void dfs(int u, int pre){
sz[u]=1;
par[u]=pre;
for (int i=0;i<adj[u].size();i++){
int v=adj[u][i];
if (v==pre)continue;
dfs(v, u);
sz[u]+=sz[v];
}
}
void hld(int u, int pre){
num[u]++;
pos[u]=num[pre];
hldpar[u]=pre;
for (int i=0;i<adj[u].size();i++){
int v=adj[u][i];
if (v==par[u])continue;
if (sz[v]*2>=sz[u])hld(v, pre);
else hld(v, v);
}
}
struct mat {
int a00, a01, a10, a11;
friend mat operator * (const mat& A, const mat& B) {
mat R;
R.a00 = min(A.a00 + B.a00, A.a01 + B.a10);
R.a01 = min(A.a00 + B.a01, A.a01 + B.a11);
R.a10 = min(A.a10 + B.a00, A.a11 + B.a10);
R.a11 = min(A.a10 + B.a01, A.a11 + B.a11);
return R;
}
};
struct segtree {
int s;
vector<mat> cost;
segtree() {}
segtree(int n) {
for (s = 1; s < n; s *= 2) {}
cost = vector<mat>(2*s);
for (int i = 0; i < 2*s; i++) {
cost[i].a01 = cost[i].a10 = 1;
}
}
void update(int i, int& v0, int& v1) {
i += s;
cost[i].a00 += v0;
cost[i].a01 += v0;
cost[i].a10 += v1;
cost[i].a11 += v1;
for (i /= 2; i; i /= 2) {
cost[i] = cost[2*i] * cost[2*i+1];
}
}
void eval(int& a0, int& a1) {
a0 = min(cost[1].a00, cost[1].a01);
a1 = min(cost[1].a10, cost[1].a11);
}
} seg[MAXN];
void initialize(int n, std::vector<int> a, std::vector<int> b) {
N = n;
for (int i = 0; i < N-1; i++) {
adj[a[i]].push_back(b[i]);
adj[b[i]].push_back(a[i]);
}
for (int i = 1; i <= N; i++) st[i] = 0;
dfs(1, 0);
hld(1, 1);
for (int i = 1; i <= N; i++) {
assert((pos[i] == 0) == (hldpar[i] == i));
if (hldpar[i] == i) seg[i] = segtree(num[i]);
}
}
void update(int cur, int v0, int v1) {
while (cur) {
int top = hldpar[cur];
int cur0, cur1, nxt0, nxt1;
seg[top].eval(cur0, cur1);
seg[top].update(pos[cur], v0, v1);
seg[top].eval(nxt0, nxt1);
v0 = min(nxt0, nxt1+1) - min(cur0, cur1+1);
v1 = min(nxt0+1, nxt1) - min(cur0+1, cur1);
cur = par[top];
}
}
int query();
int cat(int v) {
st[v] = 1;
update(v, 0, N);
return query();
}
int dog(int v) {
st[v] = 2;
update(v, N, 0);
return query();
}
int neighbor(int v) {
if (st[v] == 1) update(v, 0, -N);
else if (st[v] == 2) update(v, -N, 0);
else assert(false);
st[v] = 0;
return query();
}
int query() {
int a0, a1;
seg[1].eval(a0, a1);
return min(a0, a1);
}
Compilation message
catdog.cpp: In function 'void dfs(int, int)':
catdog.cpp:11:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
11 | for (int i=0;i<adj[u].size();i++){
| ~^~~~~~~~~~~~~~
catdog.cpp: In function 'void hld(int, int)':
catdog.cpp:22:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
22 | for (int i=0;i<adj[u].size();i++){
| ~^~~~~~~~~~~~~~
catdog.cpp: At global scope:
catdog.cpp:72:7: error: 'MAXN' was not declared in this scope; did you mean 'MX'?
72 | } seg[MAXN];
| ^~~~
| MX
catdog.cpp: In function 'void initialize(int, std::vector<int>, std::vector<int>)':
catdog.cpp:75:5: error: 'N' was not declared in this scope
75 | N = n;
| ^
catdog.cpp:85:29: error: 'seg' was not declared in this scope
85 | if (hldpar[i] == i) seg[i] = segtree(num[i]);
| ^~~
catdog.cpp: In function 'void update(int, int, int)':
catdog.cpp:93:9: error: 'seg' was not declared in this scope
93 | seg[top].eval(cur0, cur1);
| ^~~
catdog.cpp: In function 'int cat(int)':
catdog.cpp:106:18: error: 'N' was not declared in this scope
106 | update(v, 0, N);
| ^
catdog.cpp: In function 'int dog(int)':
catdog.cpp:112:15: error: 'N' was not declared in this scope
112 | update(v, N, 0);
| ^
catdog.cpp: In function 'int neighbor(int)':
catdog.cpp:117:35: error: 'N' was not declared in this scope
117 | if (st[v] == 1) update(v, 0, -N);
| ^
catdog.cpp:118:37: error: 'N' was not declared in this scope
118 | else if (st[v] == 2) update(v, -N, 0);
| ^
catdog.cpp: In function 'int query()':
catdog.cpp:126:5: error: 'seg' was not declared in this scope
126 | seg[1].eval(a0, a1);
| ^~~