顺序统计树简单实现

#include <iostream>
#include <string>
#include <cstring>
#include <fstream>
#include <functional>
#include <algorithm>
#include <ctime>
#include <cmath>
#include <vector>

using namespace std;
enum COLOR { RED, BLACK };

struct Node {
	int key;
	COLOR color;
	Node* p;
	Node* left;
	Node* right;
	int size;
	Node(int k, COLOR c = RED, Node* parent = nullptr, Node* l = nullptr, Node* r = nullptr,int s = 1) :key(k), color(c), p(parent), left(l), right(r),size(s) {}
};

class Tree {
	friend ostream& operator<<(ostream& o, const Tree& t);
private:
	Node* root;
	void leftR(Node* n) {
		if (!n || !n->right) {
			return;
		}
		Node* r = n->right;
		Node* p = n->p;
		if (p) {
			if (p->right == n)
				p->right = r;
			else
				p->left = r;
		}
		else
			root = r;
		r->p = p;
		n->right = r->left;
		if (n->right)
			n->right->p = n;
		r->left = n;
		n->p = r;
		//修改size值以维护
		r->size = n->size;
		n->size = (n->left ? n->left->size : 0) + (n->right ? n->right->size: 0) + 1;
	}
	void rightR(Node* n) {
		if (!n || !n->left) {
			return;
		}
		Node* l = n->left;
		Node* p = n->p;
		if (p) {
			if (p->left == n)
				p->left = l;
			else
				p->right = l;
		}
		else
			root = l;
		l->p = p;
		n->left = l->right;
		if (n->left)
			n->left->p = n;
		l->right = n;
		n->p = l;
		//修改size值以维护
		l->size = n->size;
		n->size = (n->left ? n->left->size : 0) + (n->right ? n->right->size : 0) + 1;
	}
	void keep(Node* tokeep) {
		while (tokeep->p && tokeep->p->color == RED) {
			if (tokeep->p->p->left == tokeep->p) {//其父为左孩子
				Node* father = tokeep->p;
				Node* uncle = father->p->right;
				if (uncle && uncle->color == RED) {
					father->color = BLACK;
					uncle->color = BLACK;
					father->p->color = RED;
					tokeep = father->p;
				}
				else {
					if (tokeep == father->right) {
						leftR(father);
						tokeep = father;
						father = father->p;
					}
					father->color = BLACK;
					father->p->color = RED;
					rightR(father->p);
				}
			}
			else {
				Node* father = tokeep->p;
				Node* uncle = father->p->left;
				if (uncle && uncle->color == RED) {
					uncle->color = BLACK;
					father->color = BLACK;
					father->p->color = RED;
					tokeep = father->p;
				}
				else {
					if (tokeep == father->left) {
						rightR(father);
						tokeep = father;
						father = father->p;
					}
					father->color = BLACK;
					father->p->color = RED;
					leftR(father->p);
				}
			}
		}
		root->color = BLACK;
	}
	ostream& pr(ostream& o, Node* r) const {
		if (!r)
			return o;
		o << r->key << " ";
		pr(o, r->left);
		pr(o, r->right);
		return o;
	}
	ostream& mr(ostream& o, Node* r) const {
		if (!r)
			return o;
		mr(o, r->left);
		o << r->key << " ";
		mr(o, r->right);
		return o;
	}
	ostream& er(ostream& o, Node* r) const {
		if (!r)
			return o;
		er(o, r->left);
		er(o, r->right);
		o << r->key << " ";
		return o;
	}
	ostream& sr(ostream& o, Node* r)const {
		if (!r)
			return o;
		sr(o, r->left);
		o << r->size << " ";
		sr(o, r->right);
		return o;
	}
	Node* getKey(int key) const {
		Node* r = root;
		while (r) {
			if (r->key == key)
				break;
			if (r->key < key)
				r = r->right;
			else
				r = r->left;
		}
		if (!r)
			return nullptr;
		return r;
	}
	Node* getMin(Node* t) {
		Node* res = nullptr;
		while (t) {
			res = t;
			t = t->left;
		}
		return res;
	}
	Node* getNext(Node* t) {
		if (t && t->right) {
			return getMin(t->right);
		}
		else {
			while (t && t->p && t->p->left == t) {
				t = t->p;
			}
			if (t && t->p)
				return t->p;
			else
				return nullptr;
		}
	}
	void dkeep(Node* x, Node* px) {
		while (x != root && (!x || x->color == BLACK)) {
			if (x == px->left) {
				Node* w = px->right;
				if (w->color == RED) {
					w->color = BLACK;
					px->color = RED;
					leftR(px);
					w = px->right;
				}
				if ((!w->left || w->left->color == BLACK) && (!w->right || w->right->color == BLACK)) {
					w->color = RED;
					x = px;
					px = px->p;
				}
				else {
					if (!w->right || w->right->color == BLACK) {
						w->color = RED;
						w->left->color = BLACK;
						rightR(w);
						w = px->right;
					}
					w->color = px->color;
					px->color = BLACK;
					w->right->color = BLACK;
					leftR(px);
					x = root;
				}
			}
			else {
				Node* w = px->left;
				if (w->color == RED) {
					w->color = BLACK;
					px->color = RED;
					rightR(px);
					w = px->left;
				}
				if ((!w->left || w->left->color == BLACK) && (!w->right || w->right->color == BLACK)) {
					w->color = RED;
					x = px;
					px = px->p;
				}
				else {
					if (!w->left || w->left->color == BLACK) {
						w->right->color = BLACK;
						w->color = RED;
						leftR(w);
						w = px->left;
					}
					w->color = px->color;
					px->color = BLACK;
					w->left->color = BLACK;
					rightR(px);
					x = root;
				}

			}
		}
		x->color = BLACK;
	}
	void clear(Node* root) {
		if (!root)
			return;
		clear(root->right);
		clear(root->left);
		delete root;
	}
	int findith(Node* root,int i) const {
		int l = root->left ? root->left->size:0;
		int s = l + 1;
		if (s == i)
			return root->key;
		else if (s > i)
			return findith(root->left, i);
		else
			return findith(root->right, i - s);
	}
	int findth(Node* root, int key,int sum = 0) const {
		if (!root)
			return sum + 1;
		int l = (root->left ? root->left->size : 0)+1;
		if (root->key == key)
			return l +sum;
		if (root->key < key)
			return findth(root->right, key, sum + l);
		else
			return findth(root->left, key, sum);
	}
public:
	Tree(Node* r = nullptr) :root(r) {}
	void insert(int key) {
		Node* tr = root;
		Node* ti = nullptr;
		while (tr) {
			++(tr->size);
			ti = tr;
			if (tr->key < key)
				tr = tr->right;
			else
				tr = tr->left;
		}
		if (!ti)
			root = new Node(key, BLACK);
		else {
			Node* tokeep = new Node(key, RED, ti);
			if (ti->key < key)
				ti->right = tokeep;
			else
				ti->left = tokeep;
			keep(tokeep);
		}
	}

	bool find(int key) const {
		return getKey(key) != nullptr;
	}
	void remove(int key) {
		Node* r = getKey(key);
		int color;
		Node* x = nullptr;
		Node* px = nullptr;
		if (!r)
			return;
		color = r->color;
		if (!r->left && !r->right) {
			x = nullptr;
			px = r->p;
			if (!px) {
				root = nullptr;
//				free(r);
				delete r;
				return;
			}
			else {
				if (px->left == r) {
					px->left = x;
				}
				else {
					px->right = x;
				}
			}
		}
		else if (!r->left) {
			x = r->right;
			px = r->p;
			if (!px) {
				root = x;
			}
			else {
				if (px->right == r) {
					px->right = x;
				}
				else {
					px->left = x;
				}
			}
			x->p = px;
		}
		else if (!r->right) {
			x = r->left;
			px = r->p;
			if (!px) {
				root = x;
			}
			else {
				if (px->right == r) {
					px->right = x;
				}
				else {
					px->left = x;
				}
			}
			x->p = px;
		}
		else {
			Node* nr = getMin(r->right);	//nr->left==nullptr
			color = nr->color;				// nr->p != nullptr

			//修改size以维护size
			nr->size = r->size;

			x = nr->right;
			px = nr->p;
			if (px->left == nr) {
				px->left = x;
			}
			else {
				px->right = x;
			}
			if (x) {
				x->p = px;
			}
			if (px == r)
				px = nr;
			if (!r->p) {
				root = nr;
			}
			else if (r->p->left == r) {
				r->p->left = nr;
			}
			else {
				r->p->right = nr;
			}
			nr->p = r->p;
			nr->left = r->left;
			nr->left->p = nr;
			nr->right = r->right;
			if (nr->right)
				nr->right->p = nr;
		}
//		free(r);
		delete r;
		
		//修改以维护size值
		Node* tmppx = px;
		while (tmppx) {
			--(tmppx->size);
			tmppx = tmppx->p;
		}

		if (color == BLACK) {
			dkeep(x, px);
		}
	}
	ostream& printsize(ostream& o) const{
		mr(o,root)	<<endl;
		sr(o, root) << endl;
		return o;
	}
	void clear() {
		clear(root);
		root = nullptr;
	}
	int findith(int i) const {
		if (i <= 0)
			return INT_MIN;
		if (root->size < i)
			return INT_MAX;
		return findith(root, i);
	}
	int findth(int key) const {
		return findth(root, key);
	}

	~Tree() {
		clear();
	}

};
ostream& operator<<(ostream& o, const Tree &t) {
	t.pr(o, t.root) << endl;
	t.mr(o, t.root) << endl;
	t.er(o, t.root) << endl;
	return o;
}

int main() {
	ifstream in("D:\\input.txt");
	//istream &in=cin;
	Tree t; 
	t.insert(1);
	t.insert(2);
	t.insert(3);
	t.insert(4);
	t.insert(8);
	t.remove(8);
	t.insert(7);
	cout << t << endl;
	t.printsize(cout);
	cout << "3th\t:\t" << t.findith(2) << endl;
	cout << "6\t:\t" << t.findth(0) << endl;
	char c;
	cin >> c;

	return 0;
}