Bzoj 3224 普通平衡樹

地址 https://ac.nowcoder.com/acm/contest/1036/A

您需要寫一種數據結構（可參考題目標題），來維護一些數，其中需要提供以下操作：

插入數值 x。
刪除數值 x(若有多個相同的數，應只刪除一個)。
查詢數值 x 的排名(若有多個相同的數，應輸出最小的排名)。
查詢排名爲 x 的數值。
求數值 x 的前驅(前驅定義爲小於 x 的最大的數)。
求數值 x 的後繼(後繼定義爲大於 x 的最小的數)。
注意： 數據保證查詢的結果一定存在。

輸入格式
第一行爲 n，表示操作的個數。

接下來 n 行每行有兩個數 opt 和 x，opt 表示操作的序號(1≤opt≤6)。

輸出格式
對於操作 3,4,5,6 每行輸出一個數，表示對應答案。

數據範圍
1≤n≤100000,所有數均在 −107 到 107 內。

輸入樣例：
8
1 10
1 20
1 30
3 20
4 2
2 10
5 25
6 -1
輸出樣例：
2
20
20
20

解答
treap樹的模板題目 這裏將模板記錄

// 111235.cpp : 此文件包含 "main" 函數。程序執行將在此處開始並結束。
//

#include <iostream>
#include <algorithm>

using  namespace std;

const int N = 100010; 
int INF = 1e8;

int n;
struct Node {
	int l, r;
	int key, val;
	int cnt, size;
}tr[N];
int root, idx;

void pushup(int p) {
	tr[p].size = tr[tr[p].l].size +
			tr[tr[p].r].size + tr[p].cnt;

}

int get_node(int key) {
	tr[++idx].key = key;
	tr[idx].val = rand();
	tr[idx].cnt = tr[idx].size = 1;
	return idx;
}

void build() {
	get_node(-INF);
	get_node(INF);
	root = 1;
	tr[1].r = 2;
	pushup(root);
}

//右旋
void zig(int &p) {
	int q = tr[p].l;
	tr[p].l = tr[q].r, tr[q].r = p; p = q;
	pushup(tr[p].r);
	pushup(p);
}

//左旋
void zag(int& p) {
	int q = tr[p].r;
	tr[p].r = tr[q].l;
	tr[q].l = p; p = q;
	pushup(tr[p].l); pushup(p);
}

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

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

int get_rank_by_key(int& p, int key) {
	if (!p) return 0;
	if (tr[p].key == key) return tr[tr[p].l].size + 1;
	if (tr[p].key > key) return get_rank_by_key(tr[p].l,key);
	return tr[tr[p].l].size + tr[p].cnt + get_rank_by_key(tr[p].r,key);
	return 0;
}


int get_key_by_rank(int& p, int rank) {
	if (!p) return INF;
	if (tr[tr[p].l].size >= rank) return get_key_by_rank(tr[p].l,rank);
	if (tr[tr[p].l].size + tr[p].cnt >= rank) return tr[p].key;
	return get_key_by_rank(tr[p].r,rank-tr[tr[p].l].size-tr[p].cnt);
	return 0;
}




int get_prev(int& p, int key) {
	if (!p)  return -INF;
	if (tr[p].key >= key)return get_prev(tr[p].l,key);
	return max(tr[p].key,get_prev(tr[p].r,key));

	return 0;
}



int get_next(int& p, int key) {
	if (!p) return INF;
	if (tr[p].key <= key) return get_next(tr[p].r, key);
	return min(tr[p].key, get_next(tr[p].l,key));

	return 0;
}



int main()
{
	build();

	cin >> n;
	while (n--) {
		int opt, x;
		cin >> opt>>x;
		if (opt == 1) {
			insert(root,x);
		}
		else if (opt == 2) {
			remove(root,x);
		}
		else if (opt == 3) {
			cout << get_rank_by_key(root,x)-1<<endl;
		}
		else if (opt == 4) {
			cout << get_key_by_rank(root,x+1)<< endl;
		}
		else if (opt == 5) {
			cout << get_prev(root, x) <<  endl;
		}
		else if (opt == 6) {
			cout << get_next(root, x) << endl;
		}
	}


	return 0;
}

我的視頻題解空間