算法导论 ch14 数据结构的扩张

支持动态集合 上的动态顺序统计，扩充红黑树。

 

1. 检索具有给定排序的元素

2. 确定一个元素的秩

 

#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; int size; 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; size = 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); int osRank(TreeNode *x) { int r = x->left->size + 1; TreeNode *y = x; while(y != root) { if (y == y->p->right) { r += y->p->left->size + 1; y = y->p; } } return r; } TreeNode* osSelect(int i) { return osSelect(root, i); } void setSize() { setSize(root); } void tDisplay(); private: TreeNode* osSelect(TreeNode*, int i); void setSize(TreeNode *&); TreeNode *root; }; // find the i-th min value TreeNode* RBTree::osSelect(TreeNode*x, int i) { int r = x->left->size + 1; if (i == r) { return x; } else if (i < r) { return osSelect(root->left, i); } else { return osSelect(root->right, i-r); } } void RBTree::setSize(TreeNode *&p) { p->size = 1; if (p->left) { setSize(p->left); p->size += p->left->size; } if (p->right) { setSize(p->right); p->size += p->right->size; } } 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->size<< ") "<< 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.setSize(); t.tDisplay(); TreeNode *p = t.osSelect(3); cout << p << endl; cout << "the rank of node is " << t.osRank(p) << endl; }