WA on #8 ~ #12
应该是 insert 和 remove 的问题。
#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 1e6 + 5, INF = 1e18;
template<class Tp, Tp INF1, Tp INF2 = -INF1>
class Bst {
private:
struct Node {
int l, r, c, s, d;
Tp v;
} a[N];
int t = 0, rt = 1;
int add(Tp v) {
a[++t] = {0, 0, 1, 1, rand(), v};
return t;
}
int _find(Tp v, int p) {
if(!p) return 0;
else if(a[p].v == v) return p;
else if(a[p].v > v) return _find(v, a[p].l);
else return _find(v, a[p].r);
}
void _insert(Tp v, int &p) {
if(!p) {
p = add(v);
return;
} else if(a[p].v == v) {
a[p].c++;
update(p);
return;
} else if(a[p].v > v) {
_insert(v, a[p].l);
if(a[a[p].l].d > a[p].d) zig(p);
} else {
_insert(v, a[p].r);
if(a[a[p].r].d > a[p].d) zag(p);
}
update(p);
}
Tp _min(int p) {
if(!a[p].l) return p;
else return _min(a[p].l);
}
Tp _max(int p) {
if(!a[p].r) return p;
else return _min(a[p].r);
}
void _remove(int v, int &p) {
if(!p) return;
if(v == a[p].v) {
if(a[p].c > 1) a[p].c--;
else if(a[p].l || a[p].r) {
if(!a[p].r || a[a[p].l].d > a[a[p].r].d) {
zig(p);
_remove(v, a[p].r);
} else {
zag(p);
_remove(v, a[p].l);
}
update(p);
} else p = 0;
return;
}
if(v < a[p].v) _remove(v, a[p].l);
else _remove(v, a[p].r);
update(p);
}
void removenext(int &p) {
if(!a[p].l) p = a[p].r;
else {
removenext(a[p].l);
update(p);
}
}
void update(int p) {
a[p].s = a[a[p].l].s + a[a[p].r].s + a[p].c;
}
int _get(int k, int p) {
if(k > a[p].s) return 2;
else if(k <= a[a[p].l].s) return _get(k, a[p].l);
else if(k <= a[a[p].l].s + a[p].c) return p;
else return _get(k - a[a[p].l].s - a[p].c, a[p].r);
}
int _rank(Tp v, int p) {
if(!p) return 1;
else if(v == a[p].v) return a[a[p].l].s + 1;
else if(v < a[p].v) return _rank(v, a[p].l);
else return _rank(v, a[p].r) + a[a[p].l].s + a[p].c;
}
int next(Tp v) {
int p = rt, r = 2;
while(p) {
if(a[p].v == v) {
if(a[p].r) {
p = a[p].r;
while(a[p].l) p = a[p].l;
return p;
}
break;
}
if(a[p].v > v && a[p].v < a[r].v) r = p;
if(a[p].v > v) p = a[p].l;
else p = a[p].r;
}
return r;
}
int pre(Tp v) {
int p = rt, r = 1;
while(p) {
if(a[p].v == v) {
if(a[p].l) {
p = a[p].l;
while(a[p].r) p = a[p].r;
return p;
}
break;
}
if(a[p].v < v && a[p].v > a[r].v) r = p;
if(a[p].v > v) p = a[p].l;
else p = a[p].r;
}
return r;
}
void zig(int &p) {
int q = a[p].l;
a[p].l = a[q].r;
a[q].r = p;
p = q;
update(a[p].r);
update(p);
}
void zag(int &p) {
int q = a[p].r;
a[p].r = a[q].l;
a[q].l = p;
p = q;
update(a[p].l);
update(p);
}
public:
void init() {
add(INF2);
add(INF1);
a[1].r = 2;
update(rt);
}
int find(Tp v) {
return _find(v, rt);
}
void insert(Tp v) {
_insert(v, rt);
}
Tp min() {
return _min(rt);
}
Tp max() {
return _min(rt);
}
Tp greater(Tp v) {
return a[next(v)].v;
}
Tp less(Tp v) {
return a[pre(v)].v;
}
void remove(Tp v) {
_remove(v, rt);
}
Tp get(int k) {
return a[_get(k + 1, rt)].v;
}
int rank(Tp v) {
return _rank(v, rt) - 1;
}
int size() {
return rank(INF1) - 1;
}
};
signed main() {
Bst<int, INF> t;
t.init();
int n;
cin>>n;
while(n--) {
int op, x;
cin>>op>>x;
if(op == 1) t.insert(x);
else if(op == 2) t.remove(x);
else if(op == 3) cout<<t.rank(x)<<endl;
else if(op == 4) cout<<t.get(x)<<endl;
else if(op == 5) cout<<t.less(x)<<endl;
else cout<<t.greater(x)<<endl;
}
}