// 枚举值表⽰颜⾊ enum Colour { RED, BLACK }; // 这⾥我们默认按key/value结构实现 template<class K, class V> struct RBTreeNode { // 这⾥更新控制平衡也要加⼊parent指针 pair<K, V> _kv; RBTreeNode<K, V>* _left; RBTreeNode<K, V>* _right; RBTreeNode<K, V>* _parent; Colour _col; RBTreeNode(const pair<K, V>& kv) :_kv(kv) , _left(nullptr) , _right(nullptr) , _parent(nullptr) {} }; template<class K, class V> class RBTree { typedef RBTreeNode<K, V> Node; public: private: Node* _root = nullptr; };红黑树的插入:
// 旋转代码的实现跟AVL树是⼀样的,只是不需要更新平衡因⼦ bool Insert(const pair<K, V>& kv) { if (_root == nullptr) { _root = new Node(kv); _root->_col = BLACK; return true; } Node* parent = nullptr; Node* cur = _root; while (cur) { if (cur->_kv.first < kv.first) { parent = cur; cur = cur->_right; } else if (cur->_kv.first > kv.first) { parent = cur; cur = cur->_left; } else { return false; } } cur = new Node(kv); // 新增结点。颜⾊红⾊给红⾊ cur->_col = RED; if (parent->_kv.first < kv.first) { parent->_right = cur; } else { parent->_left = cur; } cur->_parent = parent; while (parent && parent->_col == RED) { Node* grandfather = parent->_parent; // g // p u if (parent == grandfather->_left) { Node* uncle = grandfather->_right; if (uncle && uncle->_col == RED) { // u存在且为红 -》变⾊再继续往上处理 parent->_col = uncle->_col = BLACK; grandfather->_col = RED; cur = grandfather; parent = cur->_parent; } else { // u存在且为⿊或不存在 -》旋转+变⾊ if (cur == parent->_left) { // g // p u //c //单旋 RotateR(grandfather); parent->_col = BLACK; grandfather->_col = RED; } else { // g // p u // c //双旋 RotateL(parent); RotateR(grandfather); cur->_col = BLACK; grandfather->_col = RED; } break; } } else { // g // u p Node* uncle = grandfather->_left; // 叔叔存在且为红,-》变⾊即可 if (uncle && uncle->_col == RED) { parent->_col = uncle->_col = BLACK; grandfather->_col = RED; // 继续往上处理 cur = grandfather; parent = cur->_parent; } else // 叔叔不存在,或者存在且为⿊ { // 情况⼆:叔叔不存在或者存在且为⿊ // 旋转+变⾊ // g // u p // c if (cur == parent->_right) { RotateL(grandfather); parent->_col = BLACK; grandfather->_col = RED; } else { // g // u p // c RotateR(parent); RotateL(grandfather); cur->_col = BLACK; grandfather->_col = RED; } break; } } } _root->_col = BLACK; return true; }Node* Find(const K& key) { Node* cur = _root; while (cur) { if (cur->_kv.first < key) { cur = cur->_right; } else if (cur->_kv.first > key) { cur = cur->_left; } else { return cur; } } return nullptr; }bool Check(Node* root, int blackNum, const int refNum) { if (root == nullptr) { // 前序遍历⾛到空时,意味着⼀条路径⾛完了 //cout << blackNum << endl; if (refNum != blackNum) { cout << "存在⿊⾊结点的数量不相等的路径" << endl; return false; } return true; } // 检查孩⼦不太⽅便,因为孩⼦有两个,且不⼀定存在,反过来检查⽗亲就⽅便多了 if (root->_col == RED && root->_parent->_col == RED) { cout << root->_kv.first << "存在连续的红⾊结点" << endl; return false; } if (root->_col == BLACK) { blackNum++; } return Check(root->_left, blackNum, refNum) && Check(root->_right, blackNum, refNum); } bool IsBalance() { if (_root == nullptr) return true; if (_root->_col == RED) return false; // 参考值 int refNum = 0; Node* cur = _root; while (cur) { if (cur->_col == BLACK) { ++refNum; } cur = cur->_left; } return Check(_root, 0, refNum); }