Treap rank 求条(玄关)
查看原帖
Treap rank 求条(玄关)
908424
Walrus楼主2024/4/16 21:47

rt,问题在 rank 上,要么 WA on #1#2 要么 #10#11.

#include <bits/stdc++.h>
#define int long long
//#include <bits/extc++.h>
using namespace std;
//using namespace __gnu_pbds;

const int N = 1e5 + 5;
const int inf = 0x3fffffff;

int n, q, op, x, rt, cnt;

struct Treap {
	struct Node {
		int l, r;
		int key, val;
		int cnt, siz;

#define l(p) t[p].l
#define r(p) t[p].r
#define key(p) t[p].key
#define val(p) t[p].val
#define cnt(p) t[p].cnt
#define siz(p) t[p].siz
	} t[N];

	void pushup(int p) {
		siz(p) = siz(l(p)) + siz(r(p)) + cnt(p);
	}

	int get_node(int k) {
		int p = ++cnt;
		key(p) = k;
		val(p) = rand();
		cnt(p) = siz(p) = 1;
		return p;
	}

	void build() {
		rt = get_node(-inf), r(1) = get_node(inf);
		pushup(rt);
	}

	void zig(int &p) {
		int q = l(p);
		l(p) = r(q), r(q) = p, p = q;
		pushup(r(p)), pushup(p);
	}

	void zag(int &p) {
		int q = r(p);
		r(p) = l(q), l(q) = p, p = q;
		pushup(l(p)), pushup(p);
	}

	void insert(int &p, int k) {
		if(!p) p = get_node(k);
		else if(key(p) == k) ++cnt(p);
		else if(key(p) > k) {
			insert(l(p), k);
			if(val(l(p)) > val(p)) zig(p);
		} else {
			insert(r(p), k);
			if(val(r(p)) > val(p)) zag(p);
		}
		pushup(p);
	}

	void remove(int &p, int k) {
		if(!p) return;
		if(key(p) == k) {
			if(cnt(p) > 1) --cnt(p);
			else if(l(p) || r(p)) {
				if(!r(p) || val(l(p)) > val(r(p))) zig(p), remove(r(p), k);
				else zag(p), remove(l(p), k);
			} else p = 0;
		} else if(key(p) > k) remove(l(p), k);
		else remove(r(p), k);
		pushup(p);
	}

	int rank(int p, int k) {
		if(!p) return 0;
		if(key(p) == k) return siz(l(p)) + 1;
		if(key(p) > k) return rank(l(p), k);
		return siz(l(p)) + cnt(p) + rank(r(p), k);
	}

	int value(int p, int k) {
		if(!p) return inf;
		if(siz(l(p)) >= k) return value(l(p), k);
		if(siz(l(p)) + cnt(p) >= k) return key(p);
		return value(r(p), k - siz(l(p)) - cnt(p));
	}

	int get_pre(int p, int k) {
		if(!p) return -inf;
		if(key(p) >= k) return get_pre(l(p), k);
		return max(key(p), get_pre(r(p), k));
	}

	int get_nxt(int p, int k) {
		if(!p) return inf;
		if(key(p) <= k) return get_nxt(r(p), k);
		return min(key(p), get_nxt(l(p), k));
	}

} BST;

signed main() {
	ios::sync_with_stdio(0);
	cin.tie(nullptr), cout.tie(nullptr);

	srand((unsigned)time(0));

//	freopen("std.in", "r", stdin);
//	freopen("my.out", "w", stdout);

	BST.build();
	cin >> q;
	while(q--) {
		cin >> op >> x;
		if(op == 1)
			BST.insert(rt, x);
		else if(op == 2)
			BST.remove(rt, x);
		else if(op == 3)
			cout << BST.rank(rt, x + 1) << '\n';//问题在这,答案会差1
		else if(op == 4)
			cout << BST.value(rt, x + 1) << '\n';
		else if(op == 5)
			cout << BST.get_pre(rt, x) << '\n';
		else if(op == 6)
			cout << BST.get_nxt(rt, x) << '\n';
	}
	return 0;
}
2024/4/16 21:47
加载中...