はじめに

どうも、わたしです。今回は Go言語（Golang）で 「AVL木」 を実装してみました。データ構造についてはあんまり理解が乏しいですが、自分が学んだことをアウトプットしたいので記事に書きます！！

なるべく詳しく書くつもりなので、てっとり早く実装どうやってやるんだよという方は以下のわたしが参考にした方（karaskさん）の Github をご覧ください。

この記事はまず二分探索木について話して、そのあと平衡二分探索木の話、AVL木についてと、メインの実装というようにまとめたいと思います。

動機

いまわたしは Go で AtCoder の競プロを頑張っているんですが ABC352 の D 問題に挑戦した際、わかんなーいってなって解説見たんですが「平衡二分探索木」を使いましょうってあってなんじゃそれはとなりました。調べてみると Go には平衡二分探索木が標準ライブラリでないらしい... 自分で作るか！です

ちなみに競プロは今年 2024 年の 6 月の ABC357 からはじめました。執筆時点ではレート 490 で茶色コーダーです。緑めざして D 問題をいっぱい解けるようになるべくがんばってる感じです。

二分探索木

平衡... に移る前にまず、二分探索木についてお話しします。二分探索木とは二分木の一種です。二分木とは、から話していきます。

1️⃣ 二分木とは

二分木 とは 「各ノードの子が 2 個以下」 である根付き木のことです。木というデータ構造に関しては連結した森のことであって、森とは閉路を持たないグラフのことです。二分木はこんな感じ

二分木はどのノードも 2 個までしか子をもたないので、左右の関係をつくることができます。左の子、とか、右の子、とか言ったりできるわけです。

この二分木ですが、いろんな種類があって、全二分木や完全二分木、二分探索木、二分ヒープなどがあります。二分がゲシュタルト崩壊してきた

2️⃣ 二分探索木とは

さて、その二分木の中でも 二分探索木 とは次のような性質をもつ二分木です。

どのノードも 左の子たちの各値 ≦ そのノードの値 ≦ 右の子たちの各値 が成り立つ

どのノードに注目しても、左の子たちは自身より小さく、右の子たちは自身より大きいという状況になっています。つまり、根から見たときには、左の子を根とする部分木のすべてのノードの値は根より小さく、右の子を根とする部分木のすべてのノードの値は根より大きくなります。絵に描くとこんな感じです。

たしかに根の 10 の左側はすべて 10 以下だし、右側は 10 以上だと思います。そして 8 のノードや 14 のノードを見ると、片側がなくても大きければ右に、小さければ左に置かれています。また、最小値は一番左 に、最大値は一番右 に置かれていることも重要です。これが二分探索木です！！

3️⃣ 二分探索木のメリット

この二分探索木、なにがうれしいのかというと、挿入・削除・探索 といった基本的な操作が 平均して効率的 に行えます、これがうれしい

例えば、さっきの二分探索木で値 5 が存在するか探索することを考えてみましょう。

二分探索木の性質のおかげで、根からスタートしてそのノードの値より大きいか小さいかで右か左かに移動していけばいいのです。挿入も同様に挿入する場所を順々に探していきます。

削除はやや面倒です。葉（子をもたないノード）でないノードを削除する場合は木を再構築しないといけません。10 のノードを削除しようとすると木ではなくなってしまいます。そのためもとあった 10 の位置に別のノードを持ってくる必要があります。どれをもってくればいいのかというと、右の子たちの中で最小の値 を選んでくればいいです。そうすれば、二分探索木がもたなくてはならない性質をくずさずにすみます。

上の図は一番右に最小があると書いてあるけど左のまちがいです、バカ

さあ、二分探索木、すごく効率がよさそうなデータ構造ですよね。すでにソートされていなければならないという点はありますが、挿入や削除などは添字を付け替えなければならない配列に比べて効率的です。探索も二分探索のように行えて、線形時間よりは早そうです。

もっとも効率がよいときは各ノードが 2 つの子をもつときで、左右どちらかに偏りがなくバランスが一番取れています。このときの高さは $\log n$ になります

$n$ 個要素があるとき、$1+2+2^2+2^3+\cdots+2^m=n$ における $m+1$ が木の高さであり、これは等比数列の和の公式から $2^{m+1}-1=n$ となるので $m+1=\log(n+1)$ です。$\log(n+1)$ も $\log n$ もかわんねえべ

よって、$O(\log n)$ で、挿入・削除・探索が行えます。良い

しかし、二分探索木には大きな問題点があります

4️⃣ 二分探索木のデメリット

たとえば下のような木を考えます。

これは要素数が 15 ありますが、根の左の子は高さ（自分を含めて葉までのノード数）が 2 なのに対し、右の子は高さが 11 もあって差がえぐいことになっていますね。でもこれも二分探索木です。条件を満たしているので。終わってる二分探索木です。なぜなら、6 より大きい数の挿入も削除も探索も順番に下まで降りていくことになり、本来の効率さが損なわれています。

つまり、二分探索木を構築する、もしくはしていく際、昇順や降順で値が挿入されていくと、左右片方の枝だけがのびていき、最悪の状態になってしまうのです！これは線形リストなどと変わらないことになってしまい、計算量は $O(n)$ になってしまいます。

まとめると

二分探索木	最悪計算時間	最善計算時間
挿入	$O(n)$	$O(\log n)$
削除	$O(n)$	$O(\log n)$
探索	$O(n)$	$O(\log n)$

これでは困ったです。せっかく $O(\log n)$ で平均して計算できても最悪が線形なら意味がない... そこで、この問題を解決するべく 平衡二分探索木 というものが現れました。

平衡二分探索木

二分探索木の欠点は、挿入や削除の順番によって線形と変わらない非効率な構造になってしまう点でした。

なら、であるならば！そうならないように挿入や削除のたびに バランス を整えれば良いのです。これが 平衡二分探索木 の考え方です。英語では Self-Balancing Binary Search Tree というのでまさしくバランスをととのえるわけです。

「バランスを整える」とはどういうことでしょうか。二分探索木の問題点はなんだったか、それは「終わってる二分探索木」ができてしまうことでした。そして、二分探索木の利点は $O(\log n)$ で挿入・削除・探索ができるということ。問題点を解決して利点を最大限に活かしたい。

計算時間が常に $O(\log n)$ になってほしい。そのためには木の高さは常に $\log n$ になっていてほしい、いや、そんな贅沢はいわないから $\log n$ の定数倍くらいまでには...

よって、バランスを整える とは「木の高さが $\log n$ のオーダーになるように調整をほどこす」ということになります。全体の高さを調整するということですね。なにも工夫をしてない二分探索木は挿入や削除で「終わってしまう」ことがありました。つまりは挿入や削除のたびに、この バランスを整える を行えばいいのです

さて、平衡二分探索木にはその種類が多くあります。今回の AVL木以外にも、赤黒木、ツリープなどたくさんです。今回はその中でも、もっとも古い AVL木を頑張ってやっていきたいと思います。一番理解がしやすいと思います。他はちらっとみたけど？？？だった。

AVL木

今回の主役 AVL木 です。AVL木は初めての平衡二分探索木で 1962 年に発表されました。平衡二分探索木は、バランスを整えることが大事でしたね。この AVL木はどのようにしてバランスを整えるのでしょうか。それは 木の回転 を行います。

1️⃣ 木の回転

木の回転とは何か、図で説明するのが一番わかりやすいのでご覧ください。がんばってかきました

図がわかりやすすぎて一発で理解できたと思いますが、木の回転はアンバランスなときに根を変更して均衡をもたせることをいいます。

① この図ではまず左の子たち（左の子の部分木）の高さが 4 で、右の子が 2 になっていて、つりあっていません。左の方がおもいんです。こういうときは右に回転をさせるといい感じになります。61 の根をもとに回転させます

② 回転と言ってもぐるぐる回しすぎるとよくわかんないことになっちゃうので、具体的には 「もとの根の左の子が新しく根となるような回転」 です。もとの根 61 の左の子 31 を新しく根として昇格させます。なのでほんの数十度の回転ですね

③ ただ回転させただけでは、31 の新しい根は子が 3 人になってしまうのでダメです。これを解決するには子のつけかえを行います。どうすればいいか、それは 「新しい根のもとの右の子たちを、新しい根の新しい右の子（すなわちもとの根）の左の子へと移す」 です。

回転によって何が起きたか、それは古い根が左の子を失いました。回転前 31 は古い根の左の子だったのに、この 31、61 の親になっちゃったんです。この失ってしまった左の子のとこに、新しい根の古い右の子 43 を養子としてあげちゃいます。これで OK なのは、43 の子孫たちはもともと 31 の右の子なので、すべて値は 31 より大きいです。ですがこの子たちは、古い根 61 よりは左にいる（61 の左の子 31 の子である）ので 61 よりは小さいんです。31 以上 61 以下に収まっているので、この 43 の子たちは 61 の右に付け替えていいんです。

④ こうすることでバランスが整いました。ちゃんと根の左右の部分木の高さは両方 3 になっていてつりあっていますね

なお右が重いときは左に回転させます。

これが木の回転です。高さが左右でずれていたら、重い方の反対へ回転させます。具体的には根を変える。で、子をうつしかえる。これで完ぺきです。この作業を挿入や削除の際に行えば良いんです。

2️⃣ 二重回転

いまのでよさそうですがひとつ問題があります。つぎのような場合です

さっきと同じようにやってみたら左右の高さの差が変わってないんです。こりゃまいったです。せっかくバランスをととのえようといきようように回転させてもなんの意味もないです。でも解決策があります、それは 二重回転 です！！

さっきの図のケースで何がダメだったか考えてみましょう。回転は根を変えて子をつけかえることでした。今回はつけかえる子 7 の高さが 3 になっています。ですが、3 の左の子 2 は高さが 2 でこちらの方が少ないですよね。新しく根となる 3 は左より右のが重いんです。この場合事故が起こってしまいます。

なぜか、それは右回転では、新しい根の左の子はもとと変わりません。つけかえるのは右の子の方なのでこちらが重いとつけかえ先でも当然重いままなんです。

じゃあどうするのか、問題は新しく根となるノードが右の方が重かったこと、左右が同じになれば解決するんです。左右の重さをそろえるには...そう回転をすればいいんです。新しく根になる 3 の右が重いので、この 3 を左に回転させて、それから 9 を右回転させればいいんです！考えた人あたまがよすぎる！

②の段階で左の子 7 の左右の差が①のときより 1 大きくなっていますが、ちゃんと回転後には両方つりあうようになっています。これで完ぺきです。

左回転のときにも注意は必要です。右が重くて左回転しようとしたとき、その右の左が重かったらダメなのでいちど右回転をさせておきます。

二重回転は、つけかえる方が重かったとき、逆側に一度子を回転させてから根を回転をさせる技になります。これでもう大丈夫です。どんな問題がきても美しくバランスをたもつことができます！

3️⃣ 平衡係数

さあ、回転によってバランスを整え、その方法もわかりました。でもやみくもに回転させるわけにはいきません。そこで 平衡係数 というものを使います。

挿入や削除が行われた際、根の平衡係数を計算します。平衡係数とは 「左の子の高さ - 右の子の高さ」 で計算される数値です。これを見て回転をするかしないかを判断します。具体的には -2 のとき左回転、2 のとき右回転 です。

まず、平衡係数はつまるところ 「高さの差」 です。これが大きい、つまり左右どちらかに偏っているとまずいのでした。偏りをなくすためには回転をすればよかったのでした。つまり、「高さの差をチェックしていき大きくなりすぎたら回転する」 ということをやればいいんです。で、そのしきい値となるのが 2 なんですが、なぜ 2 なのでしょうか。

実は 1 のときに回転しても意味がないことがわかります。それはつぎのことを考えてください。

右回転、つまり左のが高さが 1 大きい状況です。これで回転をしても結局右が 1 下がるだけで、右側の高さが x+1 となってしまい、意味ないのです。

では 2 のときはどうなるのか。

さっきと同じように右側が 1 下がって高さが x+1 になります。差が 2 の場合、回転後は左側の高さが x+1 になりますから高さの差は 0 で、改善しています！

もし z が x より大きかったら差が 0 にならないじょのいこ、という疑問があると思いますが、はい、おっしゃるとおりです。z は x+1 にはなることができます（これ以上増えると高さの差自体が 2 でなくなるのでこのケースに当てはまらないです）。このときには、回転後、右が高さ x+2 になるので差は 1 です。最初の状態よりは 1 改善している、と考えましょう。

高さの差が 2 であるケースは上の場合だけではないですね。二重回転が必要な場合もあります。

この場合も二重回転を行えばしっかり高さの差をなくすことができます。

x+1=z のケースは上でやったので z は小さいです。では y はどうなんだと怒られると思います。y は x と等しいか、1 小さい場合のどちらかです。それは高さの差 2 で改善を行うので 2 は以上はありえないし、x より大きくなるとそも全体の高さの差が 2 より大きくなるからです。y がいずれの場合も、回転後の高さの差は 0 ですよね？

よって、高さの差が 2 のときにバランス調整を行えば、バランスを常に保つことができるとわかると思います。平衡係数＝「左の子の高さ - 右の子の高さ」でしたから、平衡係数が 2 のときには、左が重いので右回転、-2 のときには右が重いので左回転 すればいいわけです。

実装

さあ、ここまで AVL木の説明をみてきてだいぶ理解もしてきました。ようやく本題、実装をしていきたいと思います！！！

1️⃣ ノード

実装の際にはまず何が必要か整理をしましょう。今回のデータ構造はもとを正せば「木」です。よって、ノード という構造がまず必要ですね。AVLNode としましょう。ノードは何を持っていればいいでしょう。木の中でも二分木なのだから、左の子 と 右の子 があるといいです。left と right です。これらはノードですよね。つまり再帰的な構造になります。ノードの中に左の子のノードがあってその中に左の子のノードがあって... となって最後には nil を持てばいいでしょう。あとは key を用意しましょう。ノードの値のことです。最後に忘れてはいけないのが height つまり、そのノードの高さの情報です。AVL木は高さの差を気にする必要があるので、高さの情報は持っておきましょう。

よって、ノードはつぎのようになります。

type AVLNode struct {
	key     int
	height  int
	left    *AVLNode
	right   *AVLNode
}

left と right はポインタなのに注意です。ポインタにしといたほうがいろいろと操作がしやすいです。

2️⃣ 各種メソッド

AVL木がもっていなきゃいけない操作は挿入・削除・探索・バランスを整える、です。これらをスマートに行うためにメソッドを作っておきます。

func (n *AVLNode) getHeight() int {
	if n == nil {
		return 0
	}
	return n.height
}

func (n *AVLNode) updateHeight() {
	if n == nil {
		return
	}
    if n.left.getHeight() < n.right.getHeight() {
        n.height = 1 + n.right.getHeight()
    } else {
        n.height = 1 + n.left.getHeight()
    }
}

それぞれメソッド名のとおりです。getHeight は nil だった場合に爆発しないようにしておきます（まあたぶん初期値は 0 だから大丈夫だと思うけど）。updateHeight は高さの再計算です。重要です。左の子、右の子、どちらか高い方に 1 を加えたのが高さになります。nil チェックは必ずしましょう。

3️⃣ 回転

回転の実装をしていきます。右回転 rotateRight と、左回転 rotateLeft を用意します。回転のやり方は前に書いたとおりです。

右回転なら

左の子が新しく根となる
左の子（新しい根）の右の子は、もとの根の左の子に新しく成り代わる
もとの根は、新しい根の右の子になる

これをそのまま書けばいいです。左回転ではすべて逆になります。

func (n *AVLNode) rotateRight() *AVLNode {
	if n.left == nil {
		return n
	}

    // 左の子が新しい根となる
	newRoot := n.left
    // 左の子の右の子はもとの根の左の子になる
	n.left = newRoot.right
    // もとの根は新しい根の右の子になる
	newRoot.right = n

	newRoot.updateHeight()
	n.updateHeight()
 
	return newRoot
}

func (n *AVLNode) rotateLeft() *AVLNode {
	if n.right == nil {
		return n
	}
 
	newRoot := n.right
	n.right = newRoot.left
	newRoot.left = n

	newRoot.updateHeight()
	n.updateHeight()
 
	return newRoot
}

返り値は新しい根としておきます。基本的には根を追っていくことになるので。あと、高さの更新を忘れずにしておきましょう。高さが変わるのは子が変わっている新しい根ともとの根だけなのでこの 2 つだけでいいです。当然 nil チェックはしっかりします。

4️⃣ バランスを整える

バランスを整えるメソッドを作りましょう。どうやってやるんだったっけ。それは平衡係数をみればいいんでしたね。-2 なら右回転、2 なら左回転です。ですが、二重回転をしないといけない場合もあると思います。それもしっかりチェックしましょう。二重回転は新しく根となるはずだったノードを逆に回転させるんでしたよね。

func (n *AVLNode) rebalance() *AVLNode {
	if n == nil {
		return n
	}

    // 高さを再計算しておく
	n.updateHeight()

	// 平衡係数
	balanceFactor := n.left.getHeight() - n.right.getHeight()
 
	// 右の子の方が重い
	if balanceFactor == -2 {
		// 右の子について、左の方が重ければ二重回転が必要
		if n.right.left.getHeight() > n.right.right.getHeight() {
            // 一旦右の子を右回転しておく
			n.right = n.right.rotateRight()
		}
        // 左回転させる
		return n.rotateLeft()
	} else if balanceFactor == 2 {
		if n.left.right.getHeight() > n.left.left.getHeight() {
			n.left = n.left.rotateLeft()
		}
		return n.rotateRight()
	}
	return n
}

意外とかんたんです。バランス整えが必要なかったらやらずにそのままもらったノードを返せばいいです。

5️⃣ 挿入・削除・検索

バランス整えるメソッドを作ってしまえば、もうかんたんです。

挿入は正しい位置を探して、そこに加えるんですよね。これは再帰的に定義しましょう。そして、バランスを整える。削除は右の子孫から一番小さいのを探してそれを削除するところに持ってくるんでした。そしてバランス。探索はふつうにやります。

func (n *AVLNode) add(key int) *AVLNode {
    // 新しくノードを作ってわたす
	if n == nil {
		return &AVLNode{key, 1, nil, nil}
	}

    // 挿入する場所を再帰的に探す
	if key < n.key {
		n.left = n.left.add(key)
	} else {
		n.right = n.right.add(key)
	}
 
	return n.rebalance()
}

func (n *AVLNode) remove(key int) *AVLNode {
	if n == nil {
		return nil
	}

    // 削除するノードを探す
	if key < n.key {
		n.left = n.left.remove(key)
	} else if key > n.key {
		n.right = n.right.remove(key)
	} else {
        // 左右どちらも子をもつ場合
		if n.left != nil && n.right != nil {
			// 根を削除するときは右側から最小をさがしてきて新たな根とする
			// それは削除する
			newRoot := n.right.minNode()
			n.key = newRoot.key
			n.right = n.right.remove(newRoot.key)
		} else if n.left != nil {
			n = n.left
		} else if n.right != nil {
			n = n.right
		} else {
            // 葉だったらnilにする
			n = nil
			return n
		}
	}
	return n.rebalance()
}

func (n *AVLNode) search(key int) *AVLNode {
	if n == nil {
		return nil
	}

    // key が一致するノードを探す
	if key < n.key {
		return n.left.search(key)
	} else if key > n.key {
		return n.right.search(key)
	} else {
		return n
	}
}

func (n *AVLNode) minNode() *AVLNode {
	if n.left != nil {
		return n.left.minNode()
	} else {
		return n
	}
}

こんな感じ！最小のノードを探すメソッドを別に用意しています。バランス調整は、挿入・削除する場所までの道のりで行っていることがわかります。その動線上で変化が起きているので順々にみていけばいいのです。

6️⃣ AVL木

最後に木をつくりましょう。

type AVLTree struct {
	root *AVLNode
}

func newAVLtree() *AVLTree {
	return new(AVLTree)
}

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

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

func (t *AVLTree) search(key int) *AVLNode {
	return t.root.search(key)
}

これですべてが完了しました。やった！！！

おわりに

ずいぶん長くなってしまいましたが、Go言語で AVL木を実装してみました。ヒープと違って、最大値と最小値両方を $O(\log n)$ で見つけられるのはうれしいです。これをつかったら ABC352 の D をしっかり AC できて良かった。というかなんで Go にはないんだ平衡二分探索木。てか set がないのがほんとに。なんで Go で競プロやりはじめたんだろう。

なんかこういうデータ構造を実装してみるみたいなのはちょこちょこやっていきたいなと思いました。ちゃんとアルゴリズムとデータ構造の勉強しなおさなければ、と毎回思っている。積ん読になっている。

参考文献

以下のリンクたちが参考になりました。ありがとうございます

追記

ABC370 D を解こうとしたら、lower_bound 的なものが必要だったので、追加で実装してみました

lower_bound は、ある key より大きな最大のノードを探してもってくるやつです。二分探索でよく使うやつですね。$O(\log n)$ でいけます。

また、その反対の upper_bound もつくりました。これは key より小さい最大のノードです。

// その部分木内でkeyより大きい最小のノード なければnil
func (n *AVLNode) lowerBound(key int) *AVLNode {
	if n == nil {
		return nil
	}
	if key < n.key {
		// keyが注目ノード未満なら左の部分木を調べれば良い
		// そこになければn自身がlowerbound (左部分木のmax<keyのため)
		res := n.left.lowerBound(key)
		if res != nil {
			return res
		} else {
			return n
		}
	} else if key > n.key {
		// keyが注目モードより大きければ右を調べるのみ
		return n.right.lowerBound(key)
	} else {
		return n
	}
}

// その部分木内でkeyより小さい最大のノード なければnil
func (n *AVLNode) upperBound(key int) *AVLNode {
	if n == nil {
		return nil
	}
	// lowerBoundと逆を行う
	if key > n.key {
		res := n.right.upperBound(key)
		if res != nil {
			return res
		} else {
			return n
		}
	} else if key < n.key {
		return n.left.upperBound(key)
	} else {
		return n
	}
}


// keyより大きい最小のノード なければnil
func (t *AVLTree) lowerBound(key int) *AVLNode {
	return t.root.lowerBound(key)
}

// keyより小さい最大のノード なければnil
func (t *AVLTree) upperBound(key int) *AVLNode {
	return t.root.upperBound(key)
}

更新履歴

2024.09.29 公開
2024.09.29 左右盲が発生していたので修正
2024.10.02 追記を追記