平衡二分探索木 AVL木 を Go で実装して高さが最小限に保たれることを確認する

golangalgorithm

AVL木は二分探索木の一種である平衡二分探索木の一つで、高さを最小限に保つことで操作の計算量を抑えることができる。 変更があるたびに左右の部分木のバランスを見て高さに2以上の差がある場合、回転してリバランスする。 他の平衡二分探索木としては赤黒木などがある。

リバランス

状況によって回転の方法が少し異なる。

左部分木の左や右部部分木の右の木が高くなった場合は、次のように一重回転することでバランスを取ることができる。

一方、それ以外の場合は次のように一度逆方向に回転する二重回転が必要になる。

Goでの実装

Go で簡単な AVL木 を実装した。

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)
	}
}

ランダムなデータをリバランスあり/なしで投入して高さを比較する。

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())
	}
}

リバランスすることで高さが最小限に抑えられていることが確認できる。

* enableRebalance = false: left height = 22, right height = 16
* enableRebalance = true: left height = 12, right height = 13

参考

「アルゴリズムとデータ構造」資料10 AVL木