// Splay Tree Implementation // used for SPOJ Problem: ORDERSET #include<stdio.h> #include<string.h> #include<math.h> #include<ctype.h> #include<assert.h> #include<stdlib.h> #include<time.h> #include<vector> #include<queue> #include<stack> #include<map> #include<set> #include<iostream> #include<algorithm> #include<string> using namespace std; #define FOR(i,n) for(int i=0;i<(n);++i) #define REP(i,a,b) for(int i=(a);i<=(b);++i) #define CLR(a,x) memset(a,(x),sizeof(a)) typedef long long LL; typedef pair<int,int> pii; #define INF (1<<30) #define KEY_TYPE int class Node { private: KEY_TYPE key; int cnt; // size of the subtree rooted here Node* left; Node* right; Node* parent; int calcTreeSize(); public: bool isLeftChild(); bool isRightChild(); bool isRoot(); KEY_TYPE getKeyVal(); int getCnt(); void Rotate(); void Zig(); void ZigZig(); void ZigZag(); bool NormalBSTInsert(Node* newNode); void Splay(); void SplayUntil(Node* until); Node(); Node(const KEY_TYPE _key); ~Node(); Node* FindNearest(const KEY_TYPE _key); bool Exists(const KEY_TYPE _key); void Insert(KEY_TYPE _key); void Insert(Node* newNode); void Delete(); void Delete(KEY_TYPE _key); void PrintInOrder(); void PrintPreOrder(); KEY_TYPE findKth(int k); int Count(KEY_TYPE _key); }; Node* ROOT; bool Node::isLeftChild() { return (this->parent!=NULL && this->parent->left == this); } bool Node::isRightChild() { return (this->parent!=NULL && this->parent->right == this); } bool Node::isRoot() { return this->parent == NULL; } KEY_TYPE Node::getKeyVal() { return this->key; } int Node::getCnt() { return this->cnt; } int Node::calcTreeSize() { return 1 + (this->left?this->left->cnt:0) + (this->right?this->right->cnt:0); } void Node::Rotate() { if(isLeftChild()) { bool parWasLeftChild = this->parent?this->parent->isLeftChild():false; parent->left = this->right; if(this->right) this->right->parent = this->parent; Node* newParent = this->parent->parent; this->right = this->parent; if(this->right) this->right->parent = this; this->parent = newParent; if(this->parent != NULL) { if(parWasLeftChild) this->parent->left = this; else this->parent->right = this; } this->right->cnt = this->right->calcTreeSize(); this->cnt = this->calcTreeSize(); } else if(isRightChild()) { bool parWasLeftChild = this->parent?this->parent->isLeftChild():false; parent->right = this->left; if(this->left) this->left->parent = parent; Node* newParent = this->parent->parent; this->left = this->parent; if(this->left) this->left->parent = this; this->parent = newParent; if(this->parent != NULL) { if(parWasLeftChild) this->parent->left = this; else this->parent->right = this; } this->left->cnt = this->left->calcTreeSize(); this->cnt = this->calcTreeSize(); } } void Node::Zig() { Rotate(); } void Node::ZigZig() { this->parent->Rotate(); this->Rotate(); } void Node::ZigZag() { this->Rotate(); this->Rotate(); } void Node::Splay() { while( !this->isRoot() ) { if(this->parent->isRoot()) // Zig { this->Zig(); } else if(this->isLeftChild() == this->parent->isLeftChild()) // ZigZig { this->ZigZig(); } else // ZigZag { this->ZigZag(); } } ROOT = this; } Node::Node() { key = 0; cnt = 0; left = NULL; right = NULL; parent = NULL; } Node::Node(const KEY_TYPE _key) { key = _key; cnt = 1; left = NULL; right = NULL; parent = NULL; } Node::~Node() { } Node* Node::FindNearest(const KEY_TYPE _key) { Node *ret = this; while(true) { if(ret->key < _key) { if(ret->right) { ret = ret->right; } else break; } else if(ret->key > _key) { if(ret->left) { ret = ret->left; } else break; } else { break; } } return ret; } bool Node::Exists(const KEY_TYPE _key) { Node* temp = this->FindNearest(_key); temp->Splay(); return ROOT->key == _key; } // returns false if a node with key = newNode->key already exists in the tree bool Node::NormalBSTInsert(Node* newNode) { Node* curNode = this; bool alreadyThere = false; while(true) { if(newNode->key > curNode->key) { if(curNode->right) { curNode = curNode->right; } else { curNode->right = newNode; newNode->parent = curNode; break; } } else if(newNode->key < curNode->key) { if(curNode->left) { curNode = curNode->left; } else { curNode->left = newNode; newNode->parent = curNode; break; } } else { alreadyThere = true; break; } } if(!alreadyThere) { Node* temp = newNode->parent; while(temp) { temp->cnt = temp->calcTreeSize(); temp = temp->parent; } return true; } else return false; } void Node::Insert(KEY_TYPE _key) { this->Insert(new Node(_key)); } void Node::Insert(Node* newNode) { if(newNode == NULL) return; if(this->NormalBSTInsert(newNode)) newNode->Splay(); } void Node::SplayUntil(Node* until) { Node* grandParent = until->parent; while(this->parent != grandParent) { if(this->parent == until) // zig { this->Zig(); } else if(this->isLeftChild() == this->parent->isLeftChild()) // zigzig { this->ZigZig(); } else this->ZigZag(); } } void Node::Delete() { this->Splay(); if(this->left) { Node* maxNode = this->left->FindNearest(INF); maxNode->SplayUntil(this->left); this->left->right = this->right; if(this->right) this->right->parent = this->left; ROOT = this->left; ROOT->parent = NULL; ROOT->cnt = ROOT->calcTreeSize(); } else if(this->right) { Node* minNode = this->right->FindNearest(-INF); minNode->SplayUntil(this->right); this->right->left = this->left; if(this->left) this->left->parent = this->right; ROOT = this->right; ROOT->parent = NULL; ROOT->cnt = ROOT->calcTreeSize(); } else { ROOT = NULL; } } void Node::Delete(KEY_TYPE _key) { if(this->Exists(_key)) { Node* forDelete = ROOT; // the Exists method splays the node to the root ROOT->Delete(); delete forDelete; forDelete = NULL; } } KEY_TYPE Node::findKth(int k) { if(this->cnt < k) return -INF; Node* temp = this; while(true) { int c = (temp->left?temp->left->cnt:0) + 1; if(c == k) break; if(c < k) { k -= c; temp = temp->right; } else { temp = temp->left; } } temp->Splay(); return ROOT->key; } int Node::Count(KEY_TYPE _key) { Node* temp = this; int ret = 0; while(true) { if(temp->key < _key) { int c = (temp->left?temp->left->cnt:0)+1; ret += c; if(temp->right) { temp = temp->right; } else { break; } } else if(temp->key > _key) { if(temp->left) { temp = temp->left; } else break; } else { ret += (temp->left?temp->left->cnt:0); break; } } temp->Splay(); return ret; } void Node::PrintInOrder() // not needed for the problem { if(this->left) this->left->PrintInOrder(); cout<<this->getKeyVal()<<" "; if(this->right) this->right->PrintInOrder(); } void Node::PrintPreOrder() // not needed for the problem { cout<<"["<<this->getKeyVal()<<":"<<this->cnt<<"]"; printf("("); if(this->left) this->left->PrintPreOrder(); printf(")"); printf("("); if(this->right) this->right->PrintPreOrder(); printf(")"); } void SplayTreeInsert(KEY_TYPE key) { if(ROOT) ROOT->Insert(key); else ROOT = new Node(key); } int main() { int Q; scanf("%d",&Q); ROOT = NULL; while(Q--) { char t[5]; int p; scanf("%s%d",t,&p); if(t[0]=='I') { SplayTreeInsert(p); } else if(t[0]=='D') { if(ROOT) ROOT->Delete(p); } else if(t[0]=='K') { int v = -INF; if(ROOT) v = ROOT->findKth(p); if(v > -INF) { printf("%d\n",v); } else puts("invalid"); } else { int v = 0; if(ROOT) v = ROOT->Count(p); printf("%d\n",v); } } return 0; }
Friday, January 27, 2012
A needlessly long but Correct Splay Tree implementation
Labels:
Balanced Binary Search Tree,
Binary Search Tree,
Data Structure,
Order Statistics,
Splay Tree,
SPOJ
Subscribe to:
Post Comments (Atom)
No comments:
Post a Comment