Lilliput Steps

小さな一歩から着実に. 数学やプログラミングのことを書きます.

Left-Leaning 赤黒木

平衡2分探索木, いつか実装したいなあと思っていたので書いてみた. 疲れた.

手元のアルゴリズムイントロダクションとか、ネットの資料を色々漁って書いた. やべえ.

2-3-4木という木を2分探索木で表現すると赤黒木になるらしい.
ただ、普通の赤黒木だと実装が結構辛そうだったので, Left-leaning 赤黒木を書いてみた.

次はTreapを書いてみたい.

#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

typedef int Key;
typedef int Value;
const bool RED = true, BLACK = false;
const int null = -(1 << 30);

struct Node {
	
	Key key;
	Value val;
	Node *left, *right;
	bool color;
	
	Node (Key key, Value val, bool color) : key(key), val(val), color(color){
		left = right = (Node *)0;
	}
	
	Node() {
		left = right = (Node *)0;
	}
} *root;

class LLRB {
	
	/* function prototype declaration (必要だと思っていた)
	Value _get(Node *x, Key key);
	Node *_insert(Node *p, Key key, Value v);
	Node *_erase(Node *p, Key key);
	Key min(Node *p);
	Node *eraseMin(Node *p);
	Node *moveRedLeft(Node *p);
	Node *moveRedRight(Node *p);
	Node *fixUp(Node *p);
	Node *rotateLeft(Node *p);
	Node *rotateRight(Node *p);
	void flipColors(Node *p);
	bool isRed(Node *p);
	*/
	
	public :
		Value get(Key key){
			return (_get(root, key));
		}
		
		void insert(Key key){
			root = _insert(root, key, key);
			root->color = BLACK;
		}
		
		bool empty(){
			return (root == NULL);
		}
		
		bool find(Key key){
			return (get(key) != null);
		}
		
		void erase(Key key){
			root = _erase(root, key);
			if (root != NULL) root->color = BLACK;
		}
	
	private :
		Value _get(Node *x, Key key){
			while (x != NULL){
				if (key == x->key) return (x->val);
				else if (key < x->key) x = x->left;
				else x = x->right;
			}
			return (null);
		}
		
		Node *_insert(Node *p, Key key, Value v){
			if (p == NULL) return (new Node(key, v, RED));
			
			if (key == p->key) p->val = v;
			else if (key < p->key) p->left = _insert(p->left, key, v);
			else p->right = _insert(p->right, key, v);
			
			if (isRed(p->right) && !isRed(p->left)) p = rotateLeft(p);
			if (isRed(p->left) && isRed(p->left->left)) p = rotateRight(p);
			if (isRed(p->left) && isRed(p->right)) flipColors(p);
			
			return (p);
		}
		
		Node *_erase(Node *p, Key key){
			if (key < p->key){
				if (!isRed(p->left) && !isRed(p->left->left)){
					p = moveRedLeft(p);
				}
				p->left = _erase(p->left, key);
			}
			else {
				if (isRed(p->left)){
					p = rotateRight(p);
				}
				
				if (key == p->key && p->right == NULL){
					return (NULL);
				}
				
				if (!isRed(p->right) && !isRed(p->right->left)){
					p = moveRedRight(p);
				}
				
				if (key == p->key){
					p->val = _get(p->right, min(p->right));
					p->key = min(p->right);
					p->right = eraseMin(p->right);
				}
				else {
					p->right = _erase(p->right, key);
				}
				
			}
			return (fixUp(p));
		}
		
		Key min(Node *p){
			if (p->left == NULL) return (p->key);
			return (min(p->left));
		}
		
		Node *eraseMin(Node *p){
			if (p->left == NULL) return (NULL);
			if (!isRed(p->left) && !isRed(p->left->left)){
				p = moveRedLeft(p);
			}
			p->left = eraseMin(p->left);
			
			return (fixUp(p));
		}
		
		Node *moveRedLeft(Node *p){
			flipColors(p);
			if (isRed(p->right->left)){
				p->right = rotateRight(p->right);
				p = rotateLeft(p);
				flipColors(p);
			}
			
			return (p);
		}
		
		Node *moveRedRight(Node *p){
			flipColors(p);
			if (isRed(p->left->left)){
				p = rotateRight(p);
				flipColors(p);
			}
			
			return (p);
		}
		
		Node *fixUp(Node *p){
			if (isRed(p->right)){
				p = rotateLeft(p);
			}
			if (isRed(p->left) && isRed(p->left->left)){
				p = rotateRight(p);
			}
			if(isRed(p->left) && isRed(p->right)){
				flipColors(p);
			}
			
			return (p);
		}
		
		Node *rotateLeft(Node *p){
			Node *q = p->right;
			p->right = q->left;
			q->left = p;
			q->color = p->color;
			p->color = RED;
			return (q);
		}
		
		Node *rotateRight(Node *p){
			Node *q = p->left;
			p->left = q->right;
			q->right = p;
			q->color = p->color;
			p->color = RED;
			return (q);
		}
		
		void flipColors(Node *p){
			p->color = !p->color;
			p->left->color = !p->left->color;
			p->right->color = !p->right->color;
		}
		
		bool isRed(Node *p){
			if (p == NULL) return (false);
			return (p->color == RED);
		}
};

#define TEST_NUM (16)

int main()
{
	LLRB test;
	root = new Node;
	
	for (int i = 0; i < TEST_NUM; i++){
		test.insert(i);
	}
	
	printf("%s\n", test.find(5) ? "Yes" : "No");
	printf("%s\n", test.find(32) ? "Yes" : "No");
	
	test.erase(8);
	
	printf("%s\n", test.find(8) ? "Yes" : "No");
	
	return (0);
}