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

#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;
}