算法导论 ch13 红黑树

#include <iostream> #include <stack> #include <queue> using namespace std; enum COLOR { RED, BLACK }; class TreeNode { public: int key; TreeNode *left; TreeNode *right; TreeNode *p; COLOR color; int level; TreeNode(int k = 0, TreeNode*l = NULL, TreeNode*r = NULL, TreeNode*p = NULL, COLOR c = BLACK) : key(k), left(l), right(r), p(p), color(c) { level = 0; } }; ostream& operator <<(ostream& os, TreeNode *p) { if (p) { os << "the value of this TreeNode is: "<< p->key<< endl; } else { os << "this tree node is nil"<< endl; } return os; } class RBTree { public: RBTree() { root = NULL; } void rbInsert(TreeNode *&p); void leftRotate(TreeNode *&x); void rightRotate(TreeNode *&y); void tTraversNonRecursiveLayerOrder(TreeNode *p); void tDisplay(); private: TreeNode *root; }; void RBTree::rbInsert(TreeNode *&z) { cout << ">> enter inserting value " << z->key << endl; TreeNode *y = NULL; TreeNode *x = root; cout << "finding insert location " << endl; while (x != NULL) { y = x; if (z->key < x->key) { x = x->left; } else { x = x->right; } } z->p = y; if (y == NULL) { root = z; } else if (z->key < y->key) { y->left = z; } else { y->right = z; } z->left = NULL; z->right = NULL; z->color = RED; //rbInsertFixup(z); while ((z->p != NULL) && (z->p->color == RED)) { if (z->p == z->p->p->left) { y = z->p->p->right; if (y->color == RED) { z->p->color = BLACK; y->color = BLACK; z = z->p->p; } else if (z == z->p->right) { z = z->p; leftRotate(z); z->p->color = BLACK; z->p->p->color = RED; rightRotate(z->p->p); } else { z = z->p; rightRotate(z); z->p->color = BLACK; z->p->p->color = RED; leftRotate(z->p->p); } } } root->color = BLACK; cout << "<< exit inserting value " << z->key << endl; } void RBTree::leftRotate(TreeNode *&x) { TreeNode *y = x->right; x->right = y->left; if (y->left != NULL) { y->left->p = x; } y->p = x->p; if (x->p == NULL) { this->root = y; } else { if (x == x->p->left) { x->p->left = y; } else { x->p->right = y; } } y->left = x; x->p = y; } void RBTree::rightRotate(TreeNode *&y) { TreeNode *x = y->left; y->left = x->right; if (x->right != NULL) { x->right->p = y; } x->p = y->p; if (y->p == NULL) { this->root = x; } else { if (y == y->p->right) { y->p->right = y; } else { y->p->left = y; } } x->right = y; y->p = x; } void RBTree::tTraversNonRecursiveLayerOrder(TreeNode *p) { queue<TreeNode*> q; q.push(p); while (!q.empty()) { TreeNode *t = q.front(); q.pop(); cout << "(" << t->key << ", "<< t->color << ", " << t->level << ") " << endl; if (t->left) { cout << "left is " << t->left->key << endl; t->left->level = t->level + 1; q.push(t->left); } if (t->right) { cout << "right is " << t->right->key << endl; t->right->level = t->level + 1; q.push(t->right); } } } void RBTree::tDisplay() { tTraversNonRecursiveLayerOrder(root); } int main() { RBTree t; int a[] = { 3, 5, 2, 1, 7 }; int size = sizeof(a)/sizeof(int); for(int i = 0; i < size; i++) { TreeNode*p = new TreeNode(a[i]); t.rbInsert(p); } t.tDisplay(); }