代码

#include <cstdio>
#include <cctype>
#include <climits>
#include <iostream>
#include <algorithm>

struct AVL_Tree
{
	static AVL_Tree* Null;
	
	AVL_Tree *ch[2];
	int value, size, height;
	
	AVL_Tree(const int v = 0) : value(v), size(1), height(0)
	{
		static bool init(false);
		if (!init)
		{
			init = true;
			size = 0, height = -1;
			ch[0] = ch[1] = this;
			return;
		}
		ch[0] = ch[1] = Null;
	}
	
	void pushup()
	{
		size = ch[0]->size + ch[1]->size + 1;
		height = std::max(ch[0]->height, ch[1]->height) + 1;
	}
};
AVL_Tree *AVL_Tree::Null(new AVL_Tree);

void rotate(AVL_Tree *&o, const bool d)
{ 
	AVL_Tree *c(o->ch[!d]);
	o->ch[!d] = c->ch[d], c->ch[d] = o;
	o->pushup(), c->pushup();
	o = c;
}

void rebalance(AVL_Tree *&o, const bool d)
{
	if (o->ch[d]->height - o->ch[!d]->height > 1)
	{
		if (o->ch[d]->ch[!d]->height - o->ch[d]->ch[d]->height == 1)
		{
			rotate(o->ch[d], d);
		}
		rotate(o, !d);
	}
}

void insert(AVL_Tree *&o, const int x)
{
	if (o == AVL_Tree::Null)
	{
		o = new AVL_Tree(x);
		return;
	}
	bool d(o->value < x);
	!d ? insert(o->ch[0], x) : insert(o->ch[1], x);
	o->pushup(), rebalance(o, d);
}

int erase(AVL_Tree *&o)
{
	int res;
	if (o->ch[0] == AVL_Tree::Null)
	{
		res = o->value;
		if (o->ch[1] == AVL_Tree::Null)
			o = o->ch[0];
		else
			o->value = erase(o->ch[1]), o->pushup(), rebalance(o, 0);
		return res;
	}
	res = erase(o->ch[0]);
	o->pushup(), rebalance(o, 1);
	return res;
}


int erase(AVL_Tree *&o, const int v)
{
	int res;
	if (o->value == v)
	{
		res = o->value;
		if (o->ch[1] == AVL_Tree::Null)
			o = o->ch[0];
		else
			o->value = erase(o->ch[1]), o->pushup(), rebalance(o, 0);
		return res;
	}
	bool d(o->value < v);
	res = !d ? erase(o->ch[0], v) : erase(o->ch[1], v);
	o->pushup(), rebalance(o, !d);
	return res;
}


int at(AVL_Tree *&o, const int k)
{
	if (k == o->ch[0]->size + 1)
		return o->value;
	return k <= o->ch[0]->size ? at(o->ch[0], k) : at(o->ch[1], k - o->ch[0]->size - 1);
}

int find_by_order(AVL_Tree* o, int k) {
	if (k == o->ch[0]->size + 1)
		return o->value;
	if (k <= o->ch[0]->size)
		return find_by_order(o->ch[0], k);
	return find_by_order(o->ch[1], k - o->ch[0]->size - 1);
}

int order_of_key(AVL_Tree* o, int w) {
	if (o == AVL_Tree::Null)
		return 1;
	if (w <= o->value)
		return order_of_key(o->ch[0], w);
	else
		return order_of_key(o->ch[1], w) + o->ch[0]->size + 1;
}

int precursor(AVL_Tree* o, int w) {
	if (o == AVL_Tree::Null)
		return -INT_MAX;
	if (w <= o->value)
		return precursor(o->ch[0], w);
	else
		return std::max(o->value, precursor(o->ch[1], w));
}

int successor(AVL_Tree* o, int w) {
	if (o == AVL_Tree::Null)
		return INT_MAX;
	if (w >= o->value)
		return successor(o->ch[1], w);
	else
		return std::min(o->value, successor(o->ch[0], w));
}

int N;
AVL_Tree *root = AVL_Tree::Null;

int main(int argc, char** argv)
{
	N = 1000000;
	for (int opt, x, i = 1; i <= N; i++) {
		opt = rand() % 6 + 1;
		x = rand() * rand();
		switch (opt) {
			case 1:
				insert(root, x);
				break;
			case 2:
				erase(root, x);
				break;
			case 3:
				printf("%d\n", order_of_key(root, x));
				break;
			case 4:
				printf("%d\n", find_by_order(root, x));
				break;
			case 5:
				printf("%d\n", precursor(root, x));
				break;
			case 6:
				printf("%d\n", successor(root, x));
				break;
        }
	}
	
	return 0;
}

评测结果