学步园

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

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

2. 确定一个元素的秩

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