C++实现二叉搜索树和AVL树

/*
 * BST.h
 *
 *  Created on: Oct 31, 2015
 *      Author: chris
 */

#ifndef BST_H_
#define BST_H_

#include<iostream>

typedef int KeyType;
struct ElemType{
	int val;
	ElemType(): val(0) {}
};

typedef struct BSTNode{
	KeyType key;
	int bf;
	BSTNode *lchild, *rchild;
	ElemType data;
	BSTNode(): key(0), bf(0),
			lchild(NULL), rchild(NULL) {}
}*BST;

//BST using key
bool BSTCreateLinearOrder(BST & T, KeyType* keys, int low, int high);
bool BSTCreateSecondOptimalSW(BST & T, KeyType* keys, double* sw, int low, int high );
bool BSTCreateSecondOptimalW(BST & T,int num ,KeyType* keys, double * w);
void BSTDestroy(BST & T);

bool BSTInsertKey(BST & T, KeyType key, ElemType & e);
bool BSTDeleteKey(BST & T, KeyType key, ElemType & e);
bool BSTDeleteKey(BST & T, KeyType key);

BSTNode* BSTSearchKey(BST & T, KeyType key);
bool BSTSearchKey(BST & T, KeyType key, ElemType & e);
bool BSTSearchKey(BST & T, KeyType key, BSTNode*& pre , BSTNode*& cur);

void BSTDisplayInOrder(BST & T);
void BSTDisplayPreOrder(BST & T);

//AVL using key, bf
bool AVLInsertKey(BST & T, KeyType key, ElemType & e);
bool AVLDeleteKey(BST & T, KeyType key, ElemType & e);
bool AVLDeleteKey(BST & T, KeyType key);

#endif /* BST_H_ */

/*
 * BST.cpp
 *
 *  Created on: Oct 31, 2015
 *      Author: chris
 */

#include"BST.h"
#include<iostream>
#include<cmath>

using namespace std;

bool BSTCreateLinearOrder(BST & T, KeyType* keys, int low, int high)
{
	if( low > high )
		return false;
	else {
		int mid = (low + high) / 2;
		T = new BSTNode();
		if(!T) return false;
		T->key = keys[mid];

		if( (low <= mid-1 && !BSTCreateLinearOrder(T->lchild, keys, low, mid-1))
			|| (mid+1 <= high && !BSTCreateLinearOrder(T->rchild, keys, mid+1, high))) {
			if(T->lchild) delete T->lchild;
			if(T->rchild) delete T->rchild;
			delete T;
			T = NULL;
			return false;
		}
	}//endif
	return true;
}

bool BSTCreateSecondOptimalSW(BST & T, KeyType* keys, double* sw, int low, int high )
{   // sw: cumul sum of weights, sw[0] = 0
	if( low < 1 || low > high )
		return false;

	int i = low;
	int minw = abs(sw[high] - sw[low]);
	int dw = sw[high] + sw[low-1];
	//optimize.
	for(int j = low+1; j <= high; ++j)
		if(abs(dw-sw[j]-sw[j-1]) < minw) {
			i = j;
			minw = abs(dw - sw[j] - sw[j-1]);
		}

	T = new BSTNode();
	if(!T) return false;
	T->key = keys[i];

	//build subTree.
	if( (low < i && !BSTCreateSecondOptimalSW(T->lchild, keys, sw, low, i-1))
		|| (i < high && !BSTCreateSecondOptimalSW(T->rchild, keys, sw, i+1, high)) ) {
		delete T;
		T = NULL;
		return false;
	}
	return true;
}

bool BSTCreateSecondOptimalW(BST & T,int num ,KeyType* keys, double * w)
{
	if(num < 1)	return false;
	w[0] = 0;
	for(int i = 1; i <= num; ++i)
		w[i] = w[i] + w[i-1];
	return BSTCreateSecondOptimalSW(T, keys, w, 1, num);
}

void BSTDestroy(BST & T)
{
	if(!T) return;
	if(T->lchild) BSTDestroy(T->lchild);
	if(T->rchild) BSTDestroy(T->rchild);
	delete T;
	T = NULL;
	return;
}

bool BSTInsertKey(BST & T, KeyType key, ElemType & e)
{
	BSTNode *pre = NULL, *cur = NULL;
	if( !BSTSearchKey(T, key, pre, cur) ) {
		// load.
		cur = new BSTNode();
		if(!cur) return false;
		cur->key = key;
		cur->data = e;
		// insert
		if(!pre) T = cur;
		else if( key < pre->key )
			pre->lchild = cur;
		else
			pre->rchild = cur;
		return true;
	}
	return false;
}

bool BSTDeleteKey(BST & T, KeyType key, ElemType & e)
{
	if(!T)
		return false;
	else {
		if( key == T->key ) {
			//Delete BST Node while maintaining the props.
			BSTNode *q = NULL, *s = NULL;
			if( !T->rchild ) {
				q = T;
				T = T->lchild;
				delete q;
			}
			else if( !T->lchild ) {
				q = T;
				T = T->rchild;
				delete q;
			}
			else {
				q = T;
				s = q->lchild;
				while( s->rchild ) {
					q = s;
					s = s->rchild;
				}

				T->key = s->key;
				T->data = s->data;

				if(q != T) q->rchild = s->lchild;
				else q->lchild = s->lchild;
				delete s;

				return true;
			}//endif T child.
		}//endif key found
		else if( key < T->key )
			return BSTDeleteKey(T->lchild, key, e);
		else
			return BSTDeleteKey(T->rchild, key, e);
	} // else
	return true;
}

bool BSTDeleteKey(BST & T, KeyType key)
{
	ElemType e;
	return BSTDeleteKey(T, key, e);
}

BSTNode* BSTSearchKey(BST & T, KeyType key)
{
	BSTNode *pre = NULL, *cur = NULL;
	if( !BSTSearchKey(T, key, pre, cur) )
		return NULL;
	return cur;
}

bool BSTSearchKey(BST & T, KeyType key, ElemType & e)
{
	BSTNode *pre = NULL, *cur = NULL;
	if( !BSTSearchKey(T, key, pre, cur) )
		return false;
	e = cur->data;
	return true;
}

bool BSTSearchKey(BST & T, KeyType key, BSTNode*& pre , BSTNode*& cur) {
	cur = T;
	if(!T)  return false;
	else if( key == T->key )
		return true;
	else  {
		pre = T;
		if( key < T->key )
			return BSTSearchKey(T->lchild, key, pre, cur);
		else
			return BSTSearchKey(T->rchild, key, pre, cur);
	}
}

void BSTDisplayInOrder(BST & T)
{
	if(!T)
		cout << "NIL";
	else if(!T->lchild && !T->rchild)
		cout << T->key;
	else {
		cout << " { ";
		BSTDisplayInOrder(T->lchild);
		cout << " ( " << T->key << " ) ";
		BSTDisplayInOrder(T->rchild);
		cout << " } ";
	}
	cout.flush();
}

void BSTDisplayPreOrder(BST & T)
{
	if(!T)
		cout << "NIL";
	else if(!T->lchild && !T->rchild)
		cout << T->key;
	else {
		cout << " <" << T->key << ": ";
		cout << "L:{ ";
		BSTDisplayPreOrder(T->lchild);
		cout << " } - R:{ ";
		BSTDisplayPreOrder(T->rchild);
		cout << " }> ";
	}
	cout.flush();
}

//AVL
#define LH  1
#define EH  0
#define RH -1

void AVLRotateRight(BST & T) {
	BSTNode* lc = T->lchild;
	T->lchild = lc->rchild;
	lc->rchild = T;
	T = lc;
} //rot right.

void AVLRotateLeft(BST & T) {
	BSTNode* rc = T->rchild;
	T->rchild = rc->lchild;
	rc->lchild = T;
	T = rc;
} //rot left.

void AVLLeftBalance(BST & T) {
	BSTNode *lc = T->lchild, *rd = NULL;
	switch( lc->bf ) {
	case LH:
		T->bf = lc->bf = EH;
		AVLRotateRight(T);
		break;
	case RH:
		rd = lc->rchild;
		switch(rd->bf) {
		case LH: T->bf = RH; lc->bf = EH;  break;
		case EH: T->bf = lc->bf = EH;     break;
		case RH: T->bf = EH; lc->bf = LH; break;
		}//endsw
	    rd->bf = EH;
 	    AVLRotateLeft( T->lchild );
	    AVLRotateRight( T );
	}//endsw
}

void AVLRightBalance(BST & T) {
	BSTNode *rc = T->rchild, *ld = NULL;
	switch( rc->bf ) {
	case RH:
		T->bf = rc->bf = EH;
		AVLRotateLeft(T);
		break;
	case LH:
		ld = rc->lchild;
		switch(ld->bf) {
		case RH: T->bf = LH; rc->bf = EH; break;
		case EH: T->bf = rc->bf = EH;     break;
		case LH: T->bf = EH; rc->bf = RH; break;
		}//endsw
	    ld->bf = EH;
 	    AVLRotateRight( T->rchild );
	    AVLRotateLeft( T );
	}//endsw
}

bool _AVLInsertKey(BST & T, KeyType key, ElemType & e, bool & taller)
{
	if(!T) {
		T = new BSTNode();
		if(!T) return false;
		T->key = key;
		T->data = e;
		T->bf = EH;
		taller = true;
	} else {
		if( key == T->key ) {     //already exists.
			taller = false;
			return false;
		}
		if( key < T->key ) {
			if(!_AVLInsertKey(T->lchild, key, e, taller))
				return false;
			if(taller)
				switch(T->bf) {
				case LH: AVLLeftBalance( T ); taller = false; break; //**
				case EH: T->bf = LH; taller = true; break;
				case RH: T->bf = EH; taller = false; break;
				}
		}//endif
		else {
			if(!_AVLInsertKey(T->rchild, key, e, taller))
				return false;
			if(taller)
				switch(T->bf) {
				case LH: T->bf = EH; taller = false; break;
				case EH: T->bf = RH; taller = true; break;
				case RH: AVLRightBalance( T ); taller = false; break; //**
				}//endsw
		}//endelse
	}//endelse
	return true;
}//InsertAVL

bool AVLInsertKey(BST & T, KeyType key, ElemType & e)
{
	bool taller = false;
	return _AVLInsertKey(T, key, e, taller);
}

bool _AVLDeleteKey(BST & T, KeyType key, ElemType & e, bool & shorter)
{
	if(!T) return false;
	else {
		if( key == T->key ) {
			e = T->data;
			BSTNode *q = NULL, *s = NULL;
			if( !T->rchild ) {
				q = T;
				T = T->lchild;
				shorter = true;
				delete q;
			}
			else if( !T->lchild ) {
				q = T;
				T = T->rchild;
				shorter = true;
				delete q;
			}
			else {     // q is pre of s
				q = T;
				s = q->lchild;
				while( s->rchild ) {
					q = s;
					s = s->rchild;
				}

				T->key = s->key;
				T->data = s->data;

				if(q != T) q->rchild = s->lchild;
				else q->lchild = s->lchild;
				delete s;

				//update bf of q.
				if(q != T) {     // s is rchild of q
					switch (q->bf) {
					case EH: q->bf = LH; shorter = false; break;
					case LH: AVLLeftBalance(q); shorter = false;  break;
					case RH: q->bf = EH; shorter = true; break;
					}
				} else {     // s is lchild of q, q is T
					switch(T->bf) {
					case EH: T->bf = RH; shorter = false; break;
					case LH: T->bf = EH; shorter = true; break;
					case RH: AVLRightBalance(T); shorter = false;  break;
					}//endsw
				}//endelse
				return true;
			}//endif T child.
		}//endif key found
		else if( key < T->key ) {
			if(!_AVLDeleteKey(T->lchild, key, e, shorter))
				return false;
			if(shorter)
				switch(T->bf) {
				case EH: T->bf = RH; shorter = false; break;
				case LH: T->bf = EH; shorter = true; break;
				case RH: AVLRightBalance(T); shorter = false; break;
				}
		}//end lt
		else {
			if(!_AVLDeleteKey(T->rchild, key, e, shorter))
				return false;
			if(shorter)
				switch(T->bf) {
				case EH: T->bf = LH; shorter = false; break;
				case LH: AVLLeftBalance(T); shorter = false; break;
				case RH: T->bf = EH; shorter = true; break;
				}
		}//end gt
	} // else
	return true;
}

bool AVLDeleteKey(BST & T, KeyType key, ElemType & e)
{
	bool shorter = false;
	return _AVLDeleteKey(T, key, e, shorter);
}

bool AVLDeleteKey(BST & T, KeyType key)
{
	bool shorter = false;
	ElemType e;
	return _AVLDeleteKey(T, key, e, shorter);
}

/*
 * Main.cpp
 *
 *  Created on: Oct 31, 2015
 *      Author: chris
 */

#include"BST.h"
#include<iostream>

using namespace std;

//sw 9 1 2 3 4 5 6 7 8 9 1 2 4 9 12 16 20 23 28
//w 9 1 2 3 4 5 6 7 8 9 1 1 2 5 3 4 4 3 5
//w 5 1 2 3 4 5 1 30 2 29 3

int main(void) {
	int n;
	cout << "number of keys: "; cout.flush();
	cin >> n;

	KeyType * keys = new KeyType[n];
	if(!keys) return -1;

	cout << "Enter keys: "; cout.flush();
	for(int i = 0; i < n; ++i)
		cin >> keys[i];

	BST T = NULL;
	for(int i = 0; i < n; ++i) {
		ElemType e;
		AVLInsertKey(T, keys[i], e);
	}

	cout << "T: " << endl;
	BSTDisplayInOrder(T);
	cout << endl;
	BSTDisplayPreOrder(T);
	cout << endl;

	cout << "Enter the number of keys to del: "; cout.flush();
	cin >> n;
	while(n--) {
		KeyType key;
		cin >> key;
		AVLDeleteKey(T, key);

		cout << "T: " << endl;
		BSTDisplayInOrder(T);
		cout << endl;
		BSTDisplayPreOrder(T);
		cout << endl;
	}

	BSTDestroy(T);
	delete keys;
	return 0;
}

#ifdef TESTSOTREE
int main(void)
{
	int n;
	cout << "number of keys: "; cout.flush();
	cin >> n;

	KeyType * keys = new KeyType[n+1];
	double * w = new double[n+1];
	if(!keys || !w) {
		if(keys) delete keys;
		if(w) delete w;
		return -1;
	}

	cout << "Enter keys: "; cout.flush();
	keys[0] = 0;
	for(int i = 1; i <= n; ++i)
		cin >> keys[i];
	cout << "Enter w: "; cout.flush();
	w[0] = 0;
	for(int i = 1; i <= n; ++i)
		cin >> w[i];

	BST T;
	BSTCreateSecondOptimalW(T ,n ,keys, w );

	cout << "T:" << endl;
	BSTDisplayInOrder(T);
	cout << endl;
	BSTDisplayPreOrder(T);
	cout << endl;

	KeyType key;
	cout << "Enter key: "; cout.flush();
	cin >> key;
	BSTNode *pre = NULL, *cur = NULL;
	if(BSTSearchKey(T, key, pre, cur)) {
		if(pre)
			cout << key << " is the child of " << pre->key << endl;
		else cout << key << " is root." << endl;
	}
	else
		cout << "Not found." << endl;

	cout << "Enter key to delete: "; cout.flush();
	cin >> key;
	if(BSTDeleteKey(T, key)) {
		cout << "success." << endl;
		cout << "T after delete key:" << endl;
		BSTDisplayInOrder(T);
		cout << endl;
		BSTDisplayPreOrder(T);
		cout << endl;
	}
	else
		cout << "failed." << endl;

	BSTDestroy(T);
	delete keys;
	delete w;
	return 0;
}

#endif