Implement AVL tree, self-balancing binary search tree, in Go and confirm that the height is kept minimum
golangalgorithmAVL tree is one of the self-balancing binary search trees, which is a kind of binary search tree, and by keeping the height to a minimum, the computational complexity for operations can be reduced. Each time the tree is changed, if there is a height difference of 2 or more between left and right trees, they are rebalanced with rotation. Another self-balancing binary search tree example is a red-black tree.
Rebalance
The method of rotation is slightly different depending on the situation.
If the tree on the left of the left subtree or the tree on the right of the right subtree becomes taller, it can be balanced by performing a single rotation as follows.
On the other hand, in other cases, a double rotation, which rotates once in the opposite direction as follows, is required.
Implementation in Go
I implemented a simple AVL tree in Go.
package main
import (
"fmt"
"log"
"math"
"math/rand"
)
type node struct {
value int
right *node
left *node
height int
}
func NewTree() *node {
return nil
}
func (t *node) getLeft() *node {
if t == nil {
return nil
}
return t.left
}
func (t *node) getRight() *node {
if t == nil {
return nil
}
return t.right
}
func (t *node) getHeight() int {
if t == nil {
return 0
}
return t.height
}
func (t *node) updateHeight() {
if t == nil {
return
}
t.height = int(math.Max(float64(t.left.getHeight()), float64(t.right.getHeight())) + 1)
}
func (t *node) Add(value int, isRebalanceEnabled bool) *node {
if t == nil {
return &node{
value: value,
height: 1,
}
}
if t.value < value {
t.right = t.getRight().Add(value, isRebalanceEnabled)
} else {
t.left = t.getLeft().Add(value, isRebalanceEnabled)
}
t.updateHeight()
if isRebalanceEnabled {
return t.rebalance()
}
return t
}
func (t *node) isLeftHeavy() bool {
return t.getLeft().getHeight() > t.getRight().getHeight()
}
func (t *node) rebalance() *node {
isBalanceStable := math.Abs(float64(t.getLeft().getHeight()-t.getRight().getHeight())) <= 1
if isBalanceStable {
return t
}
if t.isLeftHeavy() {
if !t.getLeft().isLeftHeavy() {
t.left = t.getLeft().rotateLeft()
}
return t.rotateRight()
} else {
if t.getRight().isLeftHeavy() {
t.right = t.getRight().rotateRight()
}
return t.rotateLeft()
}
}
func (t *node) rotateRight() *node {
if t.getLeft() == nil {
return t
}
newRoot := t.getLeft()
t.left = t.getLeft().getRight()
newRoot.right = t
newRoot.updateHeight()
t.updateHeight()
return newRoot
}
func (t *node) rotateLeft() *node {
if t.getRight() == nil {
return t
}
newRoot := t.getRight()
t.right = t.getRight().getLeft()
newRoot.left = t
newRoot.updateHeight()
t.updateHeight()
return newRoot
}
func (t *node) Contains(value int) bool {
if t == nil {
return false
}
if t.value == value {
return true
}
if t.value < value {
return t.getRight().Contains(value)
} else {
return t.getLeft().Contains(value)
}
}
Inject random data with / without rebalancing and compare heights.
func main() {
values := make([]int, 0, 1000)
for i := 1; i <= 1000; i++ {
values = append(values, i)
}
rand.Seed(100)
rand.Shuffle(len(values), func(i, j int) {
values[i], values[j] = values[j], values[i]
})
for _, isRebalanceEnabled := range []bool{false, true} {
tree := NewTree()
for _, v := range values {
tree = tree.Add(v, isRebalanceEnabled)
}
for _, v := range values {
if !tree.Contains(v) {
log.Fatalf("* enableRebalance = %v: value %d is missing (bug)",
isRebalanceEnabled, v)
}
}
fmt.Printf("* enableRebalance = %v: left height = %d, right height = %d\n",
isRebalanceEnabled, tree.getLeft().getHeight(), tree.getRight().getHeight())
}
}
It can be confirmed that the height is minimized by rebalancing.
* enableRebalance = false: left height = 22, right height = 16
* enableRebalance = true: left height = 12, right height = 13