Implement an AVL Tree with Go (Part II)

Introduce the deletion of nodes in an AVL tree, and then design a test and benchmark function to check its performance and effectiveness.

Cover photo

I have discussed the basic concept of AVL tree in Implement an AVL Tree with Go. The article introduces how to search… | by Kashiwa | Medium

Before introducing the deletion of the AVL Tree, we should review how to remove a node in a binary tree.

1. If the node we want to delete has no subtree, we can directly remove it. In other words, we set the pointer that points to it as nil (see Fig.1). If no pointer is pointing to it, then the node is the root of the tree, and we simply set the tree’s root as nil.
2. If the node we wish to delete has only a left subtree, we can connect its parent with this left subtree. (see Fig. 2)
3. If the node we want to delete has a right subtree or only a right subtree, we can replace it with its successor. The successor is the node that represents the smallest node among those greater than the node that should be removed. (See Fig. 3)

Fig. 1 The node we want to delete, the red node, has no children.

Fig. 2 The node we want to delete, the red node, has only left child.

Fig. 3 The node we want to delete, the red node, has right child.

After deleting the node, to adhere to the rules of the AVL tree, we need to update the height and rebalance the path between the root node and the deleted node.

type AVLTree struct {
    root *AVLNode
}

type AVLNode struct {
    Key    int
    Val    interface{}
    left   *AVLNode
    right  *AVLNode
    height int // the height of the node
}

func (t *AVLTree) Delete(key int) {
    t.root = t.root.delete(key)
}

// n: the root node of the subtree
// key: the key of the node we want to delete it
func (n *AVLNode) delete(key int) *AVLNode {
    if n == nil {
        return nil
    }
    if key < n.Key {
        n.left = n.left.delete(key)
    } else if key > n.Key {
        n.right = n.right.delete(key)
    } else {
        if n.right != nil {
            // scenario 3, node has right child
            successor := n.right
            for successor.left != nil {
                successor = successor.left
            }
            // copy the key and value of the successor 
            // to the node should be deleted originally
            n.Key = successor.Key
            n.Val = successor.Val
            n.right = n.right.delete(successor.Key)
        } else if n.left != nil {
            // scenario 2, node with only the left child
            n = n.left
        } else {
            // node has no children
            n = nil
            return n  // return the new root
        }
    }
    // update the height of node
    // and then rebalance the node
    if n != nil {
        n.updateHeight()
        n = n.rebalance()
    }
    return n
}

Test

After implementing the delete function, we can design a series of test functions to test that the delete function is working properly.

func TestAVLTreeCorrectness(t *testing.T) {
    // Declare an int array of length 1024.
    list := make([]int, 1024)
    // set i-th element to be i
    for i := range list {
        list[i] = i
    }
    // shuffle the list
    // the 2nd argument is used to swap the list
    rand.Shuffle(len(list), func(i, j int) {
        list[i], list[j] = list[j], list[i]
    })

    // create a new avl tree
    tree := new(AVLTree)
    // insert node with key`e` and value`e`
    for _, e := range list {
        tree.Insert(NewAVLNode(e, e))
     }

    min := 0

    // tree.Do() is a method that traverses the tree 
    // in increasing order.
    // We implemented the method in the previous article
    tree.Do(func(n *AVLNode) {
        // check the tree is in increasing order
        if n.Key < min {
            t.Fail()
        }
        min = n.Key
        
        // check that each node's balance factor is in [-1, 1]
        if n.balanceFactor() > 1 || n.balanceFactor() < -1 {
            t.Fail()
        }

        // check that the Search() method is working properly
        if found := tree.Search(n.Key); found == nil || found.Key != n.Val.(int) {
            t.Fail()
        }
    })

    // shuffle the list again
    rand.Shuffle(len(list), func(i, j int) {
        list[i], list[j] = list[j], list[i]
    })
 
    for i := range list {
        // delete one of element in the tree
        tree.Delete(list[i])
        min := 0
        
        tree.Do(func(n *AVLNode) {
            // check the tree is in increasing order
            if n.Key < min {
                t.Fail()
            }
            min = n.Key
            // check that each node's balance factor is in [-1, 1]
            if n.balanceFactor() > 1 || n.balanceFactor() < -1 {
                t.Fail()
            }
        })
    }
}

We should write the above code to a file whose filename ends with _test , for example avl_test.go . To run the test function, please run the following command in the terminal.

go test --run TestAVLTreeCorrectness

If you want to print some information, use -v flag to enable verbose mode during the test.

go test -v --run TestAVLTreeCorrectness

Performance

Not only do we concern about correctness, but also about performance. We can compare the AVL tree we implemented with an existing package. In this section I will use github.com/karask/go-avltree as an example to compare its performance with ours.

Step 1. Create a simple dataset

func generateData() ([]int, []int) {
    // fix the random seed
    r := rand.New(rand.NewSource(0))
    // create an int array of length 1024
    list := make([]int, 1024)
    for i := range list {
        list[i] = i
    }
    // shuffle the list
    r.Shuffle(len(list), func(i, j int) {
        list[i], list[j] = list[j], list[i]
    })
    // create a second array that duplicatesthe first one
    list2 := make([]int, len(list))
    copy(list2, list)
    r.Shuffle(len(list2), func(i, j int) {
        list2[i], list2[j] = list2[j], list2[i]
    })
    return list, list2
}

Step2. Create a benchmark for the third-party package

import "github.com/karask/go-avltree"
// import other dependencies

func BenchmarkAVLKarask(b *testing.B) {
    for i := 0; i < b.N; i++ {
        list, list2 := generateData()
        tree := new(avltree.AVLTree)
        for _, key := range list {
            tree.Add(key, key)
        }
        for _, key := range list2 {
            tree.Remove(key)
        }
    }
}

Step3. Create a benchmark for the third-party package

func BenchmarkAVLOurs(b *testing.B) {
    for i := 0; i < b.N; i++ {
        list, list2 := generateData()
        tree := new(AVLTree)
        for _, key := range list {
            tree.Insert(NewAVLNode(key, key))
        }
        for _, key := range list2 {
            tree.Delete(key)
        }
    }
}

Step4. Run benchmarks

As in the previous section, we should write the above code to a file whose filename ends with _test. Of course, you can also write the benchmark function in the same file as the test function.

To run the test function, please execute the following command in the terminal. This command will run all test functions before executing benchmarks, and it will run the benchmark functions that start with BenchmarkAVL .

go test -bench BenchmarkAVL

If you do not want to run test funcion set -run to be none

go test -bench BenchmarkAVL -run none

We can view the result as shown below

goos: windows
goarch: amd64
pkg: avl
cpu: Intel(R) Core(TM) i5-8250U CPU @ 1.60GHz
BenchmarkAVLKarask-8        3314            313372 ns/op
BenchmarkAVLOurs-8          4759            252458 ns/op
PASS

The first column contains the function name concatenated with the number of threads. The second column represents how many times the function can be executed in one second. The last column shows the time consumed for a single execution.

See full code here: https://github.com/ksw2000/Medium/tree/main/implement-an-avl-tree-with-go

See the relative articles:

1. Implement a an AVL Tree with Go:
   https://ksw2000.medium.com/implement-an-avl-tree-with-go-49e5952389d4
2. Implement a Binary Search Tree with Go: https://medium.com/@ksw2000/implement-a-binary-search-tree-with-go-969e716cc47