红黑树的简单实现

#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;
	Node(int k, COLOR c = RED, Node* parent = nullptr, Node* l = nullptr, Node* r = nullptr) :key(k), color(c), p(parent), left(l), right(r) {}
};

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;
	}
	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;
	}
	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;
	}
	Node* getKey(int key) {
		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;
	}
public:
	Tree(Node* r = nullptr) :root(r) {}
	void insert(int key) {
		Node* tr = root;
		Node* ti = nullptr;
		while (tr) {
			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) {
		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);
				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
			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);
		if (color == BLACK) {
			dkeep(x,px);
		}
	}
};
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);
	cout << t << endl;
	t.insert(2);
	cout << t << endl;
	t.insert(3);
	cout << t << endl;
	t.insert(4);
	cout << t << endl;
	t.insert(8);
	cout << t << endl;
	t.insert(5);
	cout << t << endl;
	t.insert(6);
	cout << t << endl;
	t.insert(7);
	cout << t << endl;
	t.remove(3);
	cout << t << endl;
	char c;
	cin >> c;

	return 0;
}