// insert and delete status AVLTreeInsert(AVLTree *tree, ElementType data); status AVLTreeInsertRun(AVLNode **root, ElementType data); status AVLTreeDelete(AVLTree *tree, ElementType data); status AVLTreeDeleteRun(AVLNode **root, ElementType data);
// search status AVLTreeSearch(AVLTree *tree, ElementType data);
// get max and min status AVLTreeGetMax(AVLTree *tree, ElementType *result); status AVLTreeGetMin(AVLTree *tree, ElementType *result);
// get status intAVLTreeGetHeight(AVLTree *tree); intGetNodeHeight(AVLNode *node); intAVLTreeGetSize(AVLTree *tree); status AVLTreeIsEmpty(AVLTree *tree); ElementType GetNodeMin(AVLNode *node); ElementType GetNodeMax(AVLNode *node);
voidAVLTreeDestroy(AVLTree *tree) { if (tree == NULL) { printf("tree is NULL\n"); return; } // recursive destroy AVLTreeDestroyRun(tree->root); free(tree); }
// insert and delete status AVLTreeInsert(AVLTree *tree, ElementType data) { if (tree == NULL) { printf("tree is NULL\n"); return error; } status result = AVLTreeInsertRun(&tree->root, data); // check status if (result == true) { tree->size++; returntrue; } elseif (result == false) returnfalse; return error; }
/* * insert data into AVLTree * return true if insert successfully * return false if insert failed */ status AVLTreeInsertRun(AVLNode **root, ElementType data) { if (*root == NULL) { *root = (AVLNode *) malloc(sizeof(AVLNode)); if (*root == NULL) { printf("malloc failed\n"); return error; } (*root)->data = data; (*root)->height = 1; (*root)->left = (*root)->right = NULL; returntrue; }
if ((*root)->data == data) returnfalse; elseif ((*root)->data < data) { // insert into right subtree status result = AVLTreeInsertRun(&(*root)->right, data); if (result == false) returnfalse; elseif (result == error) return error;
int leftHeight = GetNodeHeight((*root)->left); int rightHeight = GetNodeHeight((*root)->right); // after insert, we need check if the tree is balanced // in this case, we insert into right subtree, // so rightHeight >= leftHeight if (rightHeight - leftHeight == 2) { if ((*root)->right->data < data) { // type RR LeftRotation(root); } else { // type RL RightRotation(&(*root)->right); LeftRotation(root); } } } else { // insert into left subtree status result = AVLTreeInsertRun(&(*root)->left, data); if (result == false) returnfalse; elseif (result == error) return error;
int leftHeight = GetNodeHeight((*root)->left); int rightHeight = GetNodeHeight((*root)->right); // after insert, we need check if the tree is balanced // in this case, we insert into left subtree, // so leftHeight >= rightHeight if (leftHeight - rightHeight == 2) { if ((*root)->left->data > data) { // type LL RightRotation(root); } else { // type LR LeftRotation(&(*root)->left); RightRotation(root); } } } // update height if abs(leftHeight - rightHeight) != 2 (*root)->height = max(GetNodeHeight((*root)->left), GetNodeHeight((*root)->right)) + 1; returntrue; }
status AVLTreeDelete(AVLTree *tree, ElementType data) { if (tree == NULL) { printf("tree is NULL\n"); return error; } status result = AVLTreeDeleteRun(&tree->root, data); // check status if (result == true) { tree->size--; returntrue; } elseif (result == false) returnfalse; return error; }
status AVLTreeDeleteRun(AVLNode **root, ElementType data) { if (*root == NULL) returnfalse; if ((*root)->data == data) { // find the node if ((*root)->left != NULL && (*root)->right != NULL) { // left and right subtree both exist
// find the min node in right subtree and replace the node ElementType key = GetNodeMin((*root)->right); (*root)->data = key;
// delete the min node in right subtree // In effect, the question is shifted to deleting the min node in right subtree AVLTreeDeleteRun(&(*root)->right, key);
int leftHeight = GetNodeHeight((*root)->left); int rightHeight = GetNodeHeight((*root)->right); // after delete, we need check if the tree is balanced // in this case, we delete the min node in right subtree, // so leftHeight >= rightHeight if (leftHeight - rightHeight == 2) { // check >= , not > if (GetNodeHeight((*root)->left->left) >= GetNodeHeight((*root)->left->right)) { // type LL RightRotation(root); } else { // type LR LeftRotation(&(*root)->left); RightRotation(root); } } } elseif ((*root)->left == NULL) { // left subtree is empty AVLNode *temp = *root; *root = (*root)->right; free(temp); temp = NULL; returntrue; } else { // right subtree is empty AVLNode *temp = *root; *root = (*root)->left; free(temp); temp = NULL; returntrue; } } elseif ((*root)->data < data) { // find in right subtree status result = AVLTreeDeleteRun(&(*root)->right, data); if (result == false) returnfalse; elseif (result == error) return error;
int leftHeight = GetNodeHeight((*root)->left); int rightHeight = GetNodeHeight((*root)->right); // after delete, we need check if the tree is balanced // in this case, we delete in right subtree, // so leftHeight >= rightHeight if (leftHeight - rightHeight == 2) { // check >= , not > if (GetNodeHeight((*root)->left->left) >= GetNodeHeight((*root)->left->right)) { // type LL RightRotation(root); } else { // type LR LeftRotation(&(*root)->left); RightRotation(root); } } } else { // find in left subtree status result = AVLTreeDeleteRun(&(*root)->left, data); if (result == false) returnfalse; elseif (result == error) return error;
int leftHeight = GetNodeHeight((*root)->left); int rightHeight = GetNodeHeight((*root)->right); // after delete, we need check if the tree is balanced // in this case, we delete in left subtree, // so rightHeight >= leftHeight if (rightHeight - leftHeight == 2) { // check >= , not > if (GetNodeHeight((*root)->right->right) >= GetNodeHeight((*root)->right->left)) { // type RR LeftRotation(root); } else { // type RL RightRotation(&(*root)->right); LeftRotation(root); } } }
package AVLTree import"slices" // node represents a node in the AVL tree.type node struct { value interface{} count int depth int left *node right *node } // newNode creates a new node with the given value. Count and depth are set to 1 by default.func newNode(e interface{}) (n *node) { return &node{ value: e, count: 1, depth: 1, } } // inOrder returns a slice of values in the subtree rooted at n in in-order traversal.func (n *node) inOrder() []interface{} { if n == nil { returnnil } if n.left == nil && n.right == nil { returnappend([]interface{}{}, slices.Repeat([]interface{}{n.value}, n.count)...) } returnappend( append(n.left.inOrder(), slices.Repeat([]interface{}{n.value}, n.count)...), n.right.inOrder()..., ) } // getDepth returns the depth of the subtree rooted at n.func (n *node) getDepth() int { if n == nil { return0 } return n.depth } // leftRotate performs a left rotation on the subtree rooted at n.func (n *node) leftRotate() *node { if n == nil || n.right == nil { return n } r := n.right n.right = r.left r.left = n n.depth = max(n.left.getDepth(), n.right.getDepth()) + 1 r.depth = max(r.left.getDepth(), r.right.getDepth()) + 1 return r } // rightRotate performs a right rotation on the subtree rooted at n.func (n *node) rightRotate() *node { if n == nil || n.left == nil { return n } l := n.left n.left = l.right l.right = n n.depth = max(n.left.getDepth(), n.right.getDepth()) + 1 l.depth = max(l.left.getDepth(), l.right.getDepth()) + 1 return l } // rightLeftRotate performs a right-left(RL) rotation on the subtree rooted at n. func(n *node) rightLeftRotate() *node { n.right = n.right.rightRotate() return n.leftRotate() } // leftRightRotate performs a left-right(LR) rotation on the subtree rooted at n. func(n *node) leftRightRotate() *node { n.left = n.left.leftRotate() return n.rightRotate() } // getMin returns the minimum value in the subtree rooted at n.func (n *node) getMin() *node { if n == nil { returnnil } if n.left == nil { return n } minNode := n.left for minNode.left != nil { minNode = minNode.left } return minNode } // adjust do necessary rotations to maintain the balance of the AVL tree.// Caller should update the pointer of n to the new root after calling this method. func(n *node) adjust() *node { if n == nil { returnnil } if n.right.getDepth() > n.left.getDepth()+1 { if n.right.right.getDepth() > n.right.left.getDepth() { // RR case n = n.leftRotate() } else { // RL case n = n.rightLeftRotate() } } elseif n.left.getDepth() > n.right.getDepth()+1 { if n.left.left.getDepth() > n.left.right.getDepth() { // LL case n = n.rightRotate() } else { // LR case n = n.leftRightRotate() } } return n } // insert inserts the given element e into the subtree rooted at n.// Caller should update the pointer of n to the new root after calling this method. // Return false if insert failed.func (n *node) insert(e interface{}, multi bool, cmp func(a interface{}, b interface{}) int) (m *node, add bool) { if n == nil { return newNode(e), true } diff := cmp(e, n.value) // already exists if diff == 0 { if multi { n.count++ return n, true } n.value = e // replace the value return n, false } if diff < 0 { // insert into left subtree if n.left, add = n.left.insert(e, multi, cmp); add { n = n.adjust() } } else { // insert into right subtree if n.right, add = n.right.insert(e, multi, cmp); add { n = n.adjust() } } n.depth = max(n.left.getDepth(), n.right.getDepth()) + 1 return n, add } // delete deletes the given element e from the subtree rooted at n.// Caller should update the pointer of n to the new root after calling this method. // Return false if delete failed.func (n *node) delete(e interface{}, cmp func(a interface{}, b interface{}) int) (m *node, del bool) { if n == nil { returnnil, false } diff := cmp(e, n.value) if diff < 0 { n.left, del = n.left.delete(e, cmp) } elseif diff > 0 { n.right, del = n.right.delete(e, cmp) } else { // delete the node del = true if n.count > 1 { n.count-- // decrease the count if multi is true } else { if n.left != nil && n.right != nil { minNode := n.right.getMin() n.value, n.count = minNode.value, minNode.count n.right, del = n.right.delete(minNode.value, cmp) } elseif n.left != nil { n = n.left } elseif n.right != nil { n = n.right } else { n = nil } } } if n != nil { // adjust the balance of the tree n.depth = max(n.left.getDepth(), n.right.getDepth()) + 1 n = n.adjust() } return n, del } // find finds the given element e in the subtree rooted at n.// Return the count of the element if found, otherwise return 0.// If multi is true, return the count of the element.// Otherwise, return 1 if found, otherwise return 0. func(n *node) find(e interface{}, multi bool, compare func(a interface{}, b interface{})int) (int, bool) { if n == nil { return0, false } cur := n for cur != nil { diff := compare(e, cur.value) if diff == 0 { if multi { return cur.count, true } return1, true } if diff < 0 { cur = cur.left } else { cur = cur.right } } return0, false }
package AVLTree import"iter" // AvlTree is an AVL tree implementation.type AvlTree struct { root *node size int compare func(a, b interface{})int multi bool// Whether to allow multiple values for the same key } // New creates a new empty AVL tree.func New(multi bool, cmp func(a, b interface{}) int) *AvlTree { if cmp == nil { panic("compare function cannot be nil") } return &AvlTree{ compare: cmp, multi: multi, } } // Iterate returns an iter.Seq2 that iterates over the AVL tree in ascending order.func (avl *AvlTree) Iterate() iter.Seq2[int, interface{}] { returnfunc(yield func(int, interface{})bool) { list := avl.inOrder() for i := 0; i < len(list); i++ { if !yield(i, list[i]) { break } } } } // Size returns the number of nodes in the AVL tree.func (avl *AvlTree) Size() int { if avl == nil { return0 } return avl.size } // inOrder returns a slice of the AVL tree in ascending order.func (avl *AvlTree) inOrder() (s []interface{}) { if avl == nil || avl.root == nil { return } return avl.root.inOrder() } // Clear removes all nodes from the AVL tree.func (avl *AvlTree) Clear() { if avl == nil { return } avl.root = nil avl.size = 0 } // Empty returns true if the AVL tree is empty.func (avl *AvlTree) Empty() bool { if avl == nil { returntrue } return avl.size == 0 } // Insert inserts a new node into the AVL tree.func (avl *AvlTree) Insert(e interface{}) { if avl == nil { return } if avl.Empty() { avl.root = newNode(e) avl.size++ return } var add bool if avl.root, add = avl.root.insert(e, avl.multi, avl.compare); add { avl.size++ } } // Delete deletes a node from the AVL tree.func (avl *AvlTree) Delete(e interface{}) { if avl == nil || avl.Empty() { return } if avl.size == 1 && avl.compare(avl.root.value, e) == 0 { avl.root = nil avl.size = 0 return } var del bool if avl.root, del = avl.root.delete(e, avl.compare); del { avl.size-- } } // Find finds a node in the AVL tree.func (avl *AvlTree) Find(e interface{}) (count int, exist bool) { if avl == nil || avl.Empty() { return0, false } return avl.root.find(e, avl.multi, avl.compare) }
main.go: (测试代码)
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
package main import avltree "AVLTree" funcmain() { t := avltree.New(true, func(a, b interface{})int { return a.(int) - b.(int) }) arr := []int{1, 2, 3, 4, 5, 6, 7, 7, 7, 8, 9, 10} for _, v := range arr { t.Insert(v) } t.Delete(4) t.Delete(5) for _, v := range t.Iterate() { print(v.(int), " ") } }