顺序统计树简单实现

#include <iostream>#include <string>#include <cstring>#include <fstream>#include <functional>#include <algorithm>#include <ctime>#include <cmath>#include <vector>using namespace std;enum COLOR { RED, BLACK };struct Node {int key;COLOR color;Node* p;Node* left;Node* right;int size;Node(int k, COLOR c = RED, Node* parent = nullptr, Node* l = nullptr, Node* r = nullptr,int s = 1) :key(k), color(c), p(parent), left(l), right(r),size(s) {}};class Tree {friend ostream& operator<<(ostream& o, const Tree& t);private:Node* root;void leftR(Node* n) {if (!n || !n->right) {return;}Node* r = n->right;Node* p = n->p;if (p) {if (p->right == n)p->right = r;elsep->left = r;}elseroot = r;r->p = p;n->right = r->left;if (n->right)n->right->p = n;r->left = n;n->p = r;//修改size值以维护r->size = n->size;n->size = (n->left ? n->left->size : 0) + (n->right ? n->right->size: 0) + 1;}void rightR(Node* n) {if (!n || !n->left) {return;}Node* l = n->left;Node* p = n->p;if (p) {if (p->left == n)p->left = l;elsep->right = l;}elseroot = l;l->p = p;n->left = l->right;if (n->left)n->left->p = n;l->right = n;n->p = l;//修改size值以维护l->size = n->size;n->size = (n->left ? n->left->size : 0) + (n->right ? n->right->size : 0) + 1;}void keep(Node* tokeep) {while (tokeep->p && tokeep->p->color == RED) {if (tokeep->p->p->left == tokeep->p) {//其父为左孩子Node* father = tokeep->p;Node* uncle = father->p->right;if (uncle && uncle->color == RED) {father->color = BLACK;uncle->color = BLACK;father->p->color = RED;tokeep = father->p;}else {if (tokeep == father->right) {leftR(father);tokeep = father;father = father->p;}father->color = BLACK;father->p->color = RED;rightR(father->p);}}else {Node* father = tokeep->p;Node* uncle = father->p->left;if (uncle && uncle->color == RED) {uncle->color = BLACK;father->color = BLACK;father->p->color = RED;tokeep = father->p;}else {if (tokeep == father->left) {rightR(father);tokeep = father;father = father->p;}father->color = BLACK;father->p->color = RED;leftR(father->p);}}}root->color = BLACK;}ostream& pr(ostream& o, Node* r) const {if (!r)return o;o << r->key << " ";pr(o, r->left);pr(o, r->right);return o;}ostream& mr(ostream& o, Node* r) const {if (!r)return o;mr(o, r->left);o << r->key << " ";mr(o, r->right);return o;}ostream& er(ostream& o, Node* r) const {if (!r)return o;er(o, r->left);er(o, r->right);o << r->key << " ";return o;}ostream& sr(ostream& o, Node* r)const {if (!r)return o;sr(o, r->left);o << r->size << " ";sr(o, r->right);return o;}Node* getKey(int key) const {Node* r = root;while (r) {if (r->key == key)break;if (r->key < key)r = r->right;elser = r->left;}if (!r)return nullptr;return r;}Node* getMin(Node* t) {Node* res = nullptr;while (t) {res = t;t = t->left;}return res;}Node* getNext(Node* t) {if (t && t->right) {return getMin(t->right);}else {while (t && t->p && t->p->left == t) {t = t->p;}if (t && t->p)return t->p;elsereturn nullptr;}}void dkeep(Node* x, Node* px) {while (x != root && (!x || x->color == BLACK)) {if (x == px->left) {Node* w = px->right;if (w->color == RED) {w->color = BLACK;px->color = RED;leftR(px);w = px->right;}if ((!w->left || w->left->color == BLACK) && (!w->right || w->right->color == BLACK)) {w->color = RED;x = px;px = px->p;}else {if (!w->right || w->right->color == BLACK) {w->color = RED;w->left->color = BLACK;rightR(w);w = px->right;}w->color = px->color;px->color = BLACK;w->right->color = BLACK;leftR(px);x = root;}}else {Node* w = px->left;if (w->color == RED) {w->color = BLACK;px->color = RED;rightR(px);w = px->left;}if ((!w->left || w->left->color == BLACK) && (!w->right || w->right->color == BLACK)) {w->color = RED;x = px;px = px->p;}else {if (!w->left || w->left->color == BLACK) {w->right->color = BLACK;w->color = RED;leftR(w);w = px->left;}w->color = px->color;px->color = BLACK;w->left->color = BLACK;rightR(px);x = root;}}}x->color = BLACK;}void clear(Node* root) {if (!root)return;clear(root->right);clear(root->left);delete root;}int findith(Node* root,int i) const {int l = root->left ? root->left->size:0;int s = l + 1;if (s == i)return root->key;else if (s > i)return findith(root->left, i);elsereturn findith(root->right, i - s);}int findth(Node* root, int key,int sum = 0) const {if (!root)return sum + 1;int l = (root->left ? root->left->size : 0)+1;if (root->key == key)return l +sum;if (root->key < key)return findth(root->right, key, sum + l);elsereturn findth(root->left, key, sum);}public:Tree(Node* r = nullptr) :root(r) {}void insert(int key) {Node* tr = root;Node* ti = nullptr;while (tr) {++(tr->size);ti = tr;if (tr->key < key)tr = tr->right;elsetr = tr->left;}if (!ti)root = new Node(key, BLACK);else {Node* tokeep = new Node(key, RED, ti);if (ti->key < key)ti->right = tokeep;elseti->left = tokeep;keep(tokeep);}}bool find(int key) const {return getKey(key) != nullptr;}void remove(int key) {Node* r = getKey(key);int color;Node* x = nullptr;Node* px = nullptr;if (!r)return;color = r->color;if (!r->left && !r->right) {x = nullptr;px = r->p;if (!px) {root = nullptr;//free(r);delete r;return;}else {if (px->left == r) {px->left = x;}else {px->right = x;}}}else if (!r->left) {x = r->right;px = r->p;if (!px) {root = x;}else {if (px->right == r) {px->right = x;}else {px->left = x;}}x->p = px;}else if (!r->right) {x = r->left;px = r->p;if (!px) {root = x;}else {if (px->right == r) {px->right = x;}else {px->left = x;}}x->p = px;}else {Node* nr = getMin(r->right);//nr->left==nullptrcolor = nr->color;// nr->p != nullptr//修改size以维护sizenr->size = r->size;x = nr->right;px = nr->p;if (px->left == nr) {px->left = x;}else {px->right = x;}if (x) {x->p = px;}if (px == r)px = nr;if (!r->p) {root = nr;}else if (r->p->left == r) {r->p->left = nr;}else {r->p->right = nr;}nr->p = r->p;nr->left = r->left;nr->left->p = nr;nr->right = r->right;if (nr->right)nr->right->p = nr;}//free(r);delete r;//修改以维护size值Node* tmppx = px;while (tmppx) {--(tmppx->size);tmppx = tmppx->p;}if (color == BLACK) {dkeep(x, px);}}ostream& printsize(ostream& o) const{mr(o,root)<<endl;sr(o, root) << endl;return o;}void clear() {clear(root);root = nullptr;}int findith(int i) const {if (i <= 0)return INT_MIN;if (root->size < i)return INT_MAX;return findith(root, i);}int findth(int key) const {return findth(root, key);}~Tree() {clear();}};ostream& operator<<(ostream& o, const Tree &t) {t.pr(o, t.root) << endl;t.mr(o, t.root) << endl;t.er(o, t.root) << endl;return o;}int main() {ifstream in("D:\\input.txt");//istream &in=cin;Tree t; t.insert(1);t.insert(2);t.insert(3);t.insert(4);t.insert(8);t.remove(8);t.insert(7);cout << t << endl;t.printsize(cout);cout << "3th\t:\t" << t.findith(2) << endl;cout << "6\t:\t" << t.findth(0) << endl;char c;cin >> c;return 0;}