类实现带旋 Treap 求调
查看原帖
类实现带旋 Treap 求调
714254
tder楼主2023/12/30 10:58

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;
	}
}
2023/12/30 10:58
加载中...