introduction

AVL Tree named after its inventor Adelson-Velskii and Landis is a sort of binary search tree. There is a balance factor -1,0,1 which represents the balance state on each node. For more simple implementation, we record the height of the substrees and calculate its balance factor dynamically.

Before goes to details, we have to know the basic operations of a tree: rotate left and rotate right.

For rotate right, supposing that we have a tree with R as the root node, we turn the whole tree right with 90 degree and then adjust it to a valid tree.

For rotate left, assuming that L is the root node of the tree, we will turn the whole tree left with 90 degree and finally it will be adjusted to a valid tree.

code

First, we define methods relevant to tree node here. Unexceptionally, there are rotate left and rotate right but we rename it with rotateClock and rotateReverse. In addition, the method fresh() was defined to recalculate the height of node.

/**
ree Node
*/

type Ordered interface {
    ~int | ~int8 | ~int16 | ~int32 | ~int64 | ~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr | ~float32 | ~float64 | ~string
}

type _AVLNode[K Ordered, V interface{}] struct {
    key    K               // key
    value  V               // value
    left   *_AVLNode[K, V] // left child
    right  *_AVLNode[K, V] // right child
    height int             // the height of the subtree, which starts from 1
}

/*
*
get the height of the left subtree
*/
func (this *_AVLNode[K, V]) leftHeight() int {
    if this.left != nil {
        return this.left.height
    }
    return 0
}

/*
*
get the height of the right subtree
*/
func (this *_AVLNode[K, V]) rightHeight() int {
    if this.right != nil {
        return this.right.height
    }
    return 0
}

/*
*
dynamically calculating the height of its own
*/
func (this *_AVLNode[K, V]) fresh() {
    leftHeight := this.leftHeight()
    rightHeigth := this.rightHeight()
    if leftHeight < rightHeigth {
        this.height = rightHeigth + 1
    } else {
        this.height = leftHeight + 1
    }
}

/*
*
rotate right
*/
func (this *_AVLNode[K, V]) rotateClock() *_AVLNode[K, V] {
    if this.left != nil {
        left := this.left
        this.left = left.right
        left.right = this
        this.fresh()
        left.fresh()
        return left
    }
    return nil
}

/*
*
rotate left
*/
func (this *_AVLNode[K, V]) rotateReverse() *_AVLNode[K, V] {
    if this.right != nil {
        right := this.right
        this.right = right.left
        right.left = this
        this.fresh()
        right.fresh()
        return right
    }
    return nil
}

We implement some operations and the data structure about AVL here. In essence, these methods is the same to Map. For examples:

Set: put a value with a key into tree
Get: get a value with a key from tree
Size: get the number of keys have set in tree
Delete: delete a key in tree
Keys(): get the all keys in tree (for simplicity, we omit it. readers can implement easily)

The following codes are some basis with only struct definition and a method Size() to get the size of keys. The more will come later!

type AVLMap[K Ordered, V interface{}] struct {
    size int             // the number of keys
    tree *_AVLNode[K, V] // the pointer to the root node of tree
}

func NewAVLMap[K Ordered, V interface{}]() *AVLMap[K, V] {
    return &AVLMap[K, V]{
        size: 0,
        tree: nil,
    }
}

func (this *AVLMap[K, V]) Size() int {
    return this.size
}

We use method Get() to get the value with a key in tree. If the key quring was not present in tree, status code false will be returned. The process looking up the given key is the same as the procedure of the binary search tree.

/*
*
get tree node
*/
func (this *AVLMap[K, V]) get(key K) *_AVLNode[K, V] {
    tree := this.tree
    for tree != nil {
        if tree.key == key {
            return tree
        } else if tree.key < key {
            tree = tree.right
        } else {
            tree = tree.left
        }
    }
    panic("key not found !")
}

func (this *AVLMap[K, V]) Get(key K) (v V, isOk bool) {
    defer func() {
        err := recover()
        if err != nil {
            v = *new(V)
            isOk = false
        }
    }()
    node := this.get(key)
    return node.value, true
}

We use method Set to set a value with a key into tree. Four scenarios will be considering on each node after setting a value with a new key.

LL: the left child of the current node's left child makes unbalance.
LR: the right child of the current node's left child makes unbalance.
RR: the right child of the current node's right child makes unbalance.
RL: the left child of the current node's right child makes unbalance.

Different scenarios have different method to rebalance, see codes for details. After backtracking, don't forget to recalculate the height of the node no matter whether current node was rebalanced.

func (this *AVLMap[K, V]) set(tree *_AVLNode[K, V], node *_AVLNode[K, V]) *_AVLNode[K, V] {
    // if tree is nil
    if tree == nil {
        this.size++
        return node
    }

    // if key exists，set value directly
    if tree.key == node.key {
        tree.value = node.value
        return tree
    }

    // if tree is not the leaf node, process recursively
    if tree.key < node.key {
        tree.right = this.set(tree.right, node)
        if tree.rightHeight()-tree.leftHeight() == 2 { // for RL or RR
            if node.key > tree.right.key {
                tree = tree.rotateReverse()
            } else {
                tree.right = tree.right.rotateClock()
                tree = tree.rotateReverse()
            }
        }
    } else {
        tree.left = this.set(tree.left, node)
        if tree.leftHeight()-tree.rightHeight() == 2 { // for LL or LR
            if node.key < tree.left.key {
                tree = tree.rotateClock()
            } else {
                tree.left = tree.left.rotateReverse()
                tree = tree.rotateClock()
            }
        }
    }
    tree.fresh()
    return tree
}

func (this *AVLMap[K, V]) Set(key K, value V) {
    node := &_AVLNode[K, V]{
        key:    key,
        value:  value,
        height: 1,
    }
    this.tree = this.set(this.tree, node)
}

We use the method Delete to delete a value with a key. The principles of deleting is that:

If the deleting node is a leaf in the tree, we delete it directly and rebalance the tree like Set.
Otherwise, three scenarios should be considered. (1) The deleting node has not left child (2) The deleting node has not right child (3) The deleting node has both left child and right chld For (1), we override the deleting node with its left child, and go to delete its left child. For (2), we ovrride the deleting node with its right child, and go to delete its right child. For (3), we exchange the deleting node with its precuror (or succeesor), and then continue to delete the deleting node recursively. For (1) - (3), the aim is to transform deleting an unleaf node to deleting a leaf node.

func (this *AVLMap[K, V]) delete(tree *_AVLNode[K, V], key K) *_AVLNode[K, V] {
    if tree == nil {
        return nil
    }

    if tree.key == key {
        // if tree is a leaf, delete dirrectly
        if tree.left == nil && tree.right == nil {
            this.size--
            return nil
        } else if tree.left == nil && tree.right != nil {
            tree.key = tree.right.key
            tree.value = tree.right.value
            tree.right = this.delete(tree.right, tree.right.key)
            if tree.right != nil {
                tree.right.fresh()
            }
        } else if tree.left != nil && tree.right == nil {
            tree.key = tree.left.key
            tree.value = tree.left.value
            tree.left = this.delete(tree.left, tree.left.key)
            if tree.left != nil {
                tree.left.fresh()
            }
        } else {
            // delete the precursor
            // first, find the pre node
            node := tree.left
            for node.right != nil {
                node = node.right
            }
            // exchange the value and key with the precursor
            tree.key, node.key = node.key, tree.key
            tree.value, node.value = node.value, tree.value
            // go to delete the precursor (now actually is current node's original key)
            tree.left = this.delete(tree.left, node.key)
            if tree.left != nil {
                tree.left.fresh()
            }
        }
    } else if tree.key > key {
        tree.left = this.delete(tree.left, key)
        if tree.left != nil {
            tree.left.fresh()
        }
    } else {
        tree.right = this.delete(tree.right, key)
        if tree.right != nil {
            tree.right.fresh()
        }
    }

    // decide rebalancing or not
    if tree.leftHeight()-tree.rightHeight() == 2 { // not balance and the left subtree is higher
        if tree.left.leftHeight() > tree.left.rightHeight() {
            tree = tree.rotateClock()
        } else {
            tree.left = tree.left.rotateReverse()
            tree = tree.rotateClock()
        }
    } else if tree.leftHeight()-tree.rightHeight() == -2 { // not balance and the right subtree is higher
        if tree.right.rightHeight() > tree.right.leftHeight() {
            tree = tree.rotateReverse()
        } else {
            tree.right = tree.right.rotateClock()
            tree = tree.rotateReverse()
        }
    }
    tree.fresh()
    return tree
}

func (this *AVLMap[K, V]) Delete(key K) {
    // delete recursively
    this.tree = this.delete(this.tree, key)
}

test

Test results show our implementation is corrently!

Test for Set.

func TestAVLMap(t *testing.T) {
    m := structure.NewAVLMap[int, string]()
    for i := 0; i < 100; i++ {
        m.Set((i), fmt.Sprint(i+100))
    }

    m.Set(0, fmt.Sprint("999"))

    for i := 100; i < 10000; i++ {
        m.Set((i), fmt.Sprint(i+100))
    }

    zero, _ := m.Get(0)
    if zero != "999" {
        panic("error")
    }

    for i := 1; i < 10000; i++ {
        v, ok := m.Get((i))
        if ok {
            if v != fmt.Sprint(i+100) {
                panic("error")
            }
        } else {
            fmt.Println("error !")
        }
    }

    fmt.Println("size:", m.Size())
}

Test for Delete.

func TestAVLMapDelete(t *testing.T) {
    m := structure.NewAVLMap[int, string]()
    for i := 0; i < 10000; i++ {
        m.Set((i), fmt.Sprint(i+100))
        //fmt.Println("insert: ", i)
    }

    for i := 3000; i < 7000; i++ {
        m.Delete(i)
    }

    for i := 0; i < 10000; i++ {
        v, ok := m.Get((i))
        if i >= 3000 && i < 7000 {
            if ok {
                panic("empty error!")
            }
        } else {
            if ok {
                if v != fmt.Sprint(i+100) {
                    panic("value error!")
                }
            } else {
                panic("unempty error!")
            }
        }
    }

    fmt.Println("size:", m.Size())

    for i := 3000; i < 7000; i++ {
        m.Set(i, fmt.Sprint(i+100))
    }
    for i := 0; i < 10000; i++ {
        v, ok := m.Get((i))
        if ok {
            if v != fmt.Sprint(i+100) {
                panic("value error!")
            }
        } else {
            panic("unempty error!")
        }
    }

    fmt.Println("size:", m.Size())
}