avl2.c

The existed tree implementation

/*
 * Usage:  avl2  list of integers ...
 *
 * Each integer will be checked to see if it is currently in
 * the AVL tree.  If not, it will be inserted. If so, it will
 * be deleted.  The tree starts out empty and the final tree is
 * printed (on its side) in an ASCII-art style
 */

#include <stdlib.h>
typedef int value_t;

#define LEFT 0
#define RIGHT 1
#define NEITHER -1
typedef int direction;
typedef struct node_s {
   value_t        value;
   struct node_s *next[2];
   int            longer:2;
} *node;
#define Balanced(n) (n->longer < 0)

node avl_find(node tree, value_t target)
{
	while (tree && target != tree->value) {
		direction next_step = (target > tree->value);
		tree = tree->next[next_step];
	}
	return tree;
}


/***************************************************
 * INSERTION                                       *
 ***************************************************/
void avl_rebalance_path(node path, value_t target)
{
	/* Each node in path is currently balanced.
	 * Until we find target, mark each node as longer
	 * in the direction of target because we know we have
	 * inserted target there
	 */
	while (path && target != path->value) {
		direction next_step = (target > path->value);
		path->longer = next_step;
		path = path->next[next_step];
	}
}

node avl_rotate_2(node *path_top, direction dir)
{
	node B, C, D, E;
	B = *path_top;
	D = B->next[dir];
	C = D->next[1-dir];
	E = D->next[dir];
	*path_top = D;
	D->next[1-dir] = B;
	B->next[dir] = C;
	B->longer = NEITHER;
	D->longer = NEITHER;
	return E;
}


node avl_rotate_3(node *path_top, direction dir, direction third)
{
	node B, F, D, C, E;
	B = *path_top;
	F = B->next[dir];
	D = F->next[1-dir];
	/* note: C and E can be NULL */
	C = D->next[1-dir];
	E = D->next[dir];
	*path_top = D;
	D->next[1-dir] = B;
	D->next[dir] = F;
	B->next[dir] = C;
	F->next[1-dir] = E;
	D->longer = NEITHER;

	/* assume both trees are balanced */
	B->longer = F->longer = NEITHER;

	if (third == NEITHER)
		return NULL;

	if (third == dir) {
		/* E holds the insertion so B is unbalanced */ 
		B->longer = 1-dir;
		return E;
	} else {
		/* C holds the insertion so F is unbalanced */
		F->longer = dir;
		return C;
	}
}


void avl_rebalance(node *path_top, value_t target)
{
	node path = *path_top;
	direction first, second, third;
	if (Balanced(path)) {
		avl_rebalance_path(path, target);
		return;
	}
	first = (target > path->value);
	if (path->longer != first) {
		/* took the shorter path */
		path->longer = NEITHER;
		avl_rebalance_path(path->next[first], target);
		return;
	}
	/* took the longer path, need to rotate */
	second = (target > path->next[first]->value);
	if (first == second) {
		/* just a two-point rotate */
		path = avl_rotate_2(path_top, first);
		avl_rebalance_path(path, target);
		return;
	}
	/* fine details of the 3 point rotate depend on the third step.
	 * However there may not be a third step, if the third point of the
	 * rotation is the newly inserted point.  In that case we record
	 * the third step as NEITHER
	 */
	path = path->next[first]->next[second];
	if (target == path->value) third = NEITHER;
	else third = (target > path->value);
	path = avl_rotate_3(path_top, first, third);
	avl_rebalance_path(path, target);
}


int avl_insert(node *treep, value_t target)
{
	/* insert the target into the tree, returning 1 on success or 0 if it
	 * already existed
	 */
	node tree = *treep;
	node *path_top = treep;
	while (tree && target != tree->value) {
		direction next_step = (target > tree->value);
		if (!Balanced(tree)) path_top = treep;
		treep = &tree->next[next_step];
		tree = *treep;
	}
	if (tree) return 0;
	tree = malloc(sizeof(*tree));
	tree->next[0] = tree->next[1] = NULL;
	tree->longer = NEITHER;
	tree->value = target;
	*treep = tree;
	avl_rebalance(path_top, target);
	return 1;
}

/******************************************************
 * DELETION                                           *
 *****************************************************/
void avl_rotate_2_shrink(node *path_top, direction dir)
{
	node D, B, C;

	D = *path_top;
	B = D->next[1-dir];
	C = B->next[dir];

	*path_top = B;
	B->next[dir] = D;
	D->next[1-dir] = C;

	if (Balanced(B)) {
		B->longer = dir;
		D->longer = 1-dir;
	} else {
		B->longer = NEITHER;
		D->longer = NEITHER;
	}
}

void avl_rotate_3_shrink(node *path_top, direction dir)
{
	node B, C, D, E, F;
	F = *path_top;
	B = F->next[1-dir];
	D = B->next[dir];
	C = D->next[1-dir];
	E = D->next[dir];

	*path_top = D;
	D->next[1-dir] = B;
	D->next[dir] = F;
	B->next[dir] = C;
	F->next[1-dir] = E;

	B->longer = NEITHER;
	F->longer = NEITHER;
	if (D->longer == dir)
		B->longer = 1-dir;
	if (D->longer == 1-dir)
		F->longer = dir;
	D->longer = NEITHER;
}

void avl_swap_del(node *targetp, node *treep, direction dir)
{
	node targetn = *targetp;
	node tree = *treep;

	*targetp = tree;
	*treep = tree->next[1-dir];
	tree->next[LEFT] = targetn->next[LEFT];
	tree->next[RIGHT]= targetn->next[RIGHT];
	tree->longer = targetn->longer;

	free(targetn);
}

node *avl_rebalance_del(node *treep, value_t target, node *targetp)
{
	/* each node from treep down towards target, but
	 * excluding the last, will have a subtree grow
	 * and need rebalancing
	 */
	node targetn = *targetp;

	while (1) {
		node tree = *treep;
		direction dir = (target > tree->value);

		if (tree->next[dir]==NULL)
			break;

		if (Balanced(tree))
			tree->longer = 1-dir;
		else if (tree->longer == dir)
			tree->longer = NEITHER;
		else {
			if (tree->next[1-dir]->longer == dir)
				avl_rotate_3_shrink(treep, dir);
			else
				avl_rotate_2_shrink(treep, dir);
			if (tree == targetn)
				targetp = &(*treep)->next[dir];
		}

		treep = &tree->next[dir];
	}
	return targetp;
}

int avl_delete(node *treep, value_t target)
{
	/* delete the target from the tree, returning 1 on success or 0 if
	 * it wasn't found
	 */
	node tree = *treep, targetn;
	node *path_top = treep;
	node *targetp = NULL;
	direction dir;
	
	while (tree) {
		dir = (target > tree->value);
		if (target == tree->value)
			targetp = treep;
		if (tree->next[dir] == NULL)
			break;
		if (Balanced(tree)
		    || (tree->longer == (1-dir) && Balanced(tree->next[1-dir]))
			) path_top = treep;
		treep = &tree->next[dir];
		tree = *treep;
	}
	if (!targetp)
		return 0;

	/* adjust balance, but don't lose 'targetp' */
	targetp = avl_rebalance_del(path_top, target, targetp);

	/* We have re-balanced everything, it remains only to 
	 * swap the end of the path (*treep == treep) with the deleted item
	 * (*targetp == targetn)
	 */
	avl_swap_del(targetp, treep, dir);

	return 1;
}

/****************************************************************
 * PRINTING and TESTING                                         *
 ***************************************************************/
void avl_print(node tree, char *prefix, char *x, direction side)
{
	int pl = strlen(prefix);
	if (tree) {
		char t;
		strcat(prefix, "   ");
		prefix[pl-1] = x[0];
		avl_print(tree->next[LEFT], prefix, " ,|", LEFT);
		prefix[pl-1] = x[1];
		t = prefix[pl-2];
		printf("%.*s--+= %d %c\n", pl, prefix, tree->value, ":/\\"[tree->longer+1]);
		prefix[pl-2] = t;
		prefix[pl-1] = x[2];
		avl_print(tree->next[RIGHT],prefix, "|` ", RIGHT);
		prefix[pl] = '\0';
	} else printf("%.*s\n", pl-1, prefix);
}

int dir_check_depth(node tree)
{
	if (tree) {
		int err = 0;
		int rv;
		int b = dir_check_depth(tree->next[LEFT]);
		int f = dir_check_depth(tree->next[RIGHT]);
		if (b == f) {
			if (!Balanced(tree))
				err = 1;
			rv = b+1;
		} else if (b == f-1) {
			if (tree->longer != RIGHT)
				err = 1;
			rv = f+1;
		} else if (b-1 == f) {
			if (tree->longer != LEFT)
				err = 1;
			rv = b+1;
		} else {
			err = 1;
			rv = 0;
		}
		if (err)
			printf("err at %d: b=%d f=%d longer=%d\n", tree->value,
			       b, f, tree->longer);
		return rv;
	}
	return 0;
}

int main(int argc, char *argv[])
{
	node tree = NULL;
	char prefix[100];
	int arg;
	for (arg=1; arg<argc; arg++) {	
		value_t v = atoi(argv[arg]);
		if (!avl_find(tree, v)) {
			if (avl_insert(&tree, v))
				printf("Inserted %d\n", v);
			else printf("Strange failure to insert %d\n", v);
		} else {
			if (avl_delete(&tree, v))
				printf("Deleted: %d\n", v);
			else printf("Strange failure to delete %d\n", v);
		}
		dir_check_depth(tree);
	}

	strcpy(prefix, "   ");
	avl_print(tree, prefix, " = ", NEITHER);
	exit(0);
}