平衡二叉树

洛谷P3369

#include <iostream>
#include <istream>
#include <sstream>
#include <vector>
#include <stack>
#include <list>
#include <map>
#include <set>
#include <deque>
#include <queue>
#include <cstring>
#include <unordered_map>
#include <unordered_set>
#include <algorithm>
#include <numeric>
#include <chrono>
#include <ctime>
#include <cmath>
#include <cctype>
#include <string>
#include <cstdio>
#include <iomanip>


#include <thread>
#include <mutex>
#include <condition_variable>
#include <functional>
#include <iterator>
using namespace std;

template <class _Type>
class AVL
{
public:
	template <class _T>
	struct AVLNode {
		AVLNode(_T input) :l(nullptr), r(nullptr), h(1), data(input) {}
		AVLNode() :l(nullptr), r(nullptr), h(1), c(1), sub(1) {}
		_T data;
		AVLNode* l, *r;
		int h, c, sub;//高度 副本个数 子树节点个数
	};
	typedef AVLNode<_Type>  value_Type;
	typedef AVLNode<_Type>* point;
public:
	AVL() :m_root(nullptr), m_tot(0) {}
	~AVL() { clear(m_root); }
public:
	typename AVL<_Type>::point GetNode(_Type&& input);

	int  GetHight(typename AVL<_Type>::point node) {
		if (node == nullptr) return 0;
		return node->h;
	}
	int  GetSubCnt(typename AVL<_Type>::point node) {
		if (node == nullptr) return 0;
		return node->sub;
	}
	void flush(typename AVL<_Type>::point node) {
		if (node == nullptr) return;
		node->h = max(GetHight(node->l), GetHight(node->r)) + 1;
		node->sub = GetSubCnt(node->l) + GetSubCnt(node->r) + node->c;
	}
public:
	typename AVL<_Type>::point insert(_Type input);//插入
	typename AVL<_Type>::point _insert(typename AVL<_Type>::point cur, _Type&& input);

	void remove(_Type input);//删除
	typename AVL<_Type>::point _remove(typename AVL<_Type>::point cur, _Type&& input);

	typename AVL<_Type>::point findKth(int k);	//查找第k大
	int findRank(_Type x);
	typename AVL<_Type>::point findPre(_Type x);	//查找x前驱
	typename AVL<_Type>::point findNext(_Type x);//查找x后继
	typename AVL<_Type>::point find(_Type x);//查找x

private:
	void  rotate_left(typename AVL<_Type>::point& input);	//左旋转
	void  rotate_right(typename AVL<_Type>::point& input);	//右旋转
	typename AVL<_Type>::point revise(typename AVL<_Type>::point& node);//修正
	void clear(typename AVL<_Type>::point& input);
	void print(typename AVL<_Type>::point& input) {
		if (input == nullptr) return;
		print(input->l);
		cout << input->data << " ";
		print(input->r);
	}
public:
	typename AVL<_Type>::point   m_root;
	typename AVL<_Type>::point   m_new;//最后一个新增节点
	int		m_tot;//总节点个数
};

template <class _Type>
typename AVL<_Type>::point AVL<_Type>::GetNode(_Type&& input) {
	AVL<_Type>::point ret = new AVL<_Type>::value_Type();
	if (ret)ret->data = input, m_tot++;
	return m_new = ret;
}

template <class _Type>
void AVL<_Type>::clear(typename AVL<_Type>::point& input) {
	if (input == nullptr) return;
	clear(input->l);
	clear(input->r);
	input->l = nullptr;
	input->r = nullptr;
	delete(input);
	input = nullptr;
}

template <class _Type>
typename AVL<_Type>::point AVL<_Type>::insert(_Type input)//插入
{
	m_new = nullptr;
	m_root = _insert(m_root, std::forward<_Type>(input));
	return m_new;
}

template <class _Type>
typename AVL<_Type>::point AVL<_Type>::_insert(typename AVL<_Type>::point p, _Type&& input)//插入
{
	//找到节点位置
	if (p == nullptr) return GetNode(std::forward<_Type>(input));

	//相同元素 标记数量
	if (p->data == input) {
		p->c++;
		p->sub++;
		return p;
	}
	if (p->data > input)p->l = _insert(p->l, std::forward<_Type>(input));
	else p->r = _insert(p->r, std::forward<_Type>(input));

	return revise(p);
}

template <class _Type>
typename AVL<_Type>::point AVL<_Type>::revise(typename AVL<_Type>::point& p) {
	if (p == nullptr) return p;
	int hl = GetHight(p->l), hr = GetHight(p->r), diff = hl - hr;
	if (diff > 1) {
		//LL
		if (GetHight(p->l->l) >= GetHight(p->l->r)) {
			rotate_right(p);
		}
		else {
			//LR
			rotate_left(p->l);
			rotate_right(p);
		}
	}
	else if (diff < -1) {
		//RL
		if (GetHight(p->r->l) >= GetHight(p->r->r)) {
			rotate_right(p->r);
			rotate_left(p);
		}
		else rotate_left(p);//RR
	}
	flush(p->l);
	flush(p->r);
	flush(p);
	return p;
}

template <class _Type>
void AVL<_Type>::remove(_Type input)//删除
{
	m_root = _remove(m_root, std::forward<_Type>(input));
}

template <class _Type>
typename AVL<_Type>::point AVL<_Type>::_remove(typename AVL<_Type>::point p, _Type&& input)
{
	if (p == nullptr) return p;
	if (p->data > input) {
		p->l = _remove(p->l, std::forward<_Type>(input));
	}
	else if (p->data < input) {
		p->r = _remove(p->r, std::forward<_Type>(input));
	}
	else {
		//已找到节点 有多个副本
		if (--p->c > 0) { m_tot--; return revise(p); }
		else {
			//有左右子树
			if (p->l && p->r) {
				//寻找前驱节点 一定有前驱和后继
				AVL<_Type>::point preplace = p->l;
				while (preplace->r) preplace = preplace->r;
				//交换数据
				swap(p->data, preplace->data);
				//副本数量
				swap(p->c, preplace->c);

				p->l = _remove(p->l, std::forward<_Type>(input));
			}
			else {
				m_tot--;
				AVL<_Type>::point preplace = p->l ? p->l : p->r;
				delete p;
				return preplace;
			}
		}
	}
	return revise(p);
}

template <class _Type>
void  AVL<_Type>::rotate_left(typename AVL<_Type>::point& node)//旋转
{
	if (node == nullptr || node->r == nullptr) return;
	AVL::point r = node->r;
	AVL::point rl = r->l;
	r->l = node;
	node->r = rl;
	node = r;

	flush(node->l);
	flush(node);
}

template <class _Type>
void  AVL<_Type>::rotate_right(typename AVL<_Type>::point& node)//右旋转
{
	if (node == nullptr || node->l == nullptr) return;
	AVL::point l = node->l;
	AVL::point lr = l->r;
	l->r = node;
	node->l = lr;
	node = l;
	flush(node->r);
	flush(node);
}

template <class _Type>
typename AVL<_Type>::point AVL<_Type>::findKth(int k)	//查找第k大
{
	if (m_root == nullptr) return nullptr;
	AVL<_Type>::point cur = m_root;
	while (cur) {
		int l = GetSubCnt(cur->l);
		int r = GetSubCnt(cur->r);
		int m = cur->c;
		if (k <= l) {
			cur = cur->l;
		}
		else if (k <= l + m) {
			return cur;
		}
		else {
			k -= l + m;
			cur = cur->r;
		}
	}
	return nullptr;
}

template <class _Type>
int AVL<_Type>::findRank(_Type x) {
	AVL<_Type>::point cur = m_root;

	int ans = 0;
	while (cur) {
		if (cur->data > x) {
			cur = cur->l;
		}
		else if (cur->data < x) {
			ans += GetSubCnt(cur->l) + cur->c;
			cur = cur->r;
		}
		else break;
	}
	if (cur == nullptr) return ans + 1;
	return ans + GetSubCnt(cur->l) + 1;
}


template <class _Type>
typename AVL<_Type>::point AVL<_Type>::findPre(_Type x)	//查找x前驱
{
	AVL<_Type>::point cur = m_root, ret = nullptr;
	while (cur) {
		if (cur->data >= x)cur = cur->l;
		else ret = cur, cur = cur->r;
	}
	return ret;
}

template <class _Type>
typename AVL<_Type>::point AVL<_Type>::findNext(_Type x)//查找x后继
{
	AVL<_Type>::point cur = m_root, ret = nullptr;
	while (cur) {
		if (cur->data <= x)cur = cur->r;
		else ret = cur, cur = cur->l;
	}
	return ret;
}

template <class _Type>
typename AVL<_Type>::point AVL<_Type>::find(_Type x)//查找x
{
	AVL<_Type>::point cur = m_root;
	while (cur) {
		if (cur->data == x) return cur;
		else if (cur->data > x) cur = cur->l;
		else cur = cur->r;
	}
	return nullptr;
}


inline int read(int& x) {
	char ch = getchar();
	int f = 1; x = 0;
	while (ch > '9' || ch < '0') { if (ch == '-')f = -1; ch = getchar(); }
	while (ch >= '0' && ch <= '9') { x = (x << 1) + (x << 3) + ch - '0'; ch = getchar(); }
	return x * f;
}

static auto speedup = []() {ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr); return nullptr; }();


int n;
int main(){
	AVL<int> tree;
	cin >> n;

	int a, b;
	for (int i = 1; i <= n; i++){
		cin >> a >> b;
		switch (a)
		{
		case 1:tree.insert(b);break;
		case 2:tree.remove(b);break;
		case 3:cout << tree.findRank(b) << endl;  break;
		case 4:
		{
			AVL<int>::point ret = tree.findKth(b);
			if (ret == nullptr) cout << -1 << endl;
			else cout << ret->data << endl; 
			break;
		}
		case 5:
		{
			AVL<int>::point ret = tree.findPre(b);
			if (ret == nullptr) cout << -1 << endl;
			else cout << ret->data << endl; 
			break;
		}
		case 6:
		{
			AVL<int>::point ret = tree.findNext(b);
			if (ret == nullptr) cout << -1 << endl;
			else cout << ret->data << endl;
			break;
		}
		default:
			break;
		}
	}
	return 0;
}