本記事は、ソフトバンクパブリッシングから発行されている「定本 Cプログラマのためのアルゴリズムとデータ構造 (SOFTBANK BOOKS)」を参考にPythonでアルゴリズムとデータ構造について学習していきます。
今回は、平衡木の中のAVL木について学んでいきます。
平衡木とは
前回の二分探索木では、平均的な計算量はO(log n)でありながらも最悪な場合はO(n)になってしまうことを学びました。
上記の最悪な場合が発生する要因は、根から見た左右の部分木の高さが異なっていることが原因となります。
これを改善するには、木の高さをlog2 n程度に調整する必要があります。
このような木を平衡木（Balanced Tree）と言います。
AVL木とは
AVL木とは、最初の平衡木として1962年にAdel'son、Vel'skiiおよびLandisが考案し、それぞれの頭文字をとり命名されました。
これは二分探索木に、「すべての節の左右部分木の高さの差を1以内に収める」という目標を追加したものです。
木の高さが1より大きくなった場合、以下の2つのパターンに応じて、木を変形させます。
- 追加された要素が根の左部分木（右部分木）の左部分木（右部分木）に追加された
- 追加された要素が根の左部分木（右部分木）の右部分木（左部分木）に追加された
パターン1では1重回転（Single Rotation）と呼ばれる操作で、木を平衡に保ちます。
パターン2では2重回転（Double Rotation）と呼ばれる操作で、木を平衡に保ちます。
以下は2重回転を行う場合を想定し、新たに要素15を挿入します。
上記で15を挿入したことにより、高さの差が2になってしまったので、以下のように調整をします。
PythonでAVL木を実装してみる
それでは、PythonでAVL木を実装してみます。
今回は挿入のみの操作に限定し、挿入の度に平衡具合を確認し、左右のバランスをとるようにしています。
# avl_tree.py
import math
from collections import defaultdict
class Node:
def __init__(self, key):
self.key = key
self.left = None
self.right = None
class AVLTree:
def __init__(self, key):
self.root = Node(key)
self.max_depth = 0
self.total = 1
self.nodes = defaultdict(list)
def balance(self):
return int(math.log2(self.total))
def single_left_rotate(self, new_root):
tmp = new_root.right
new_root.right = self.root
self.root.left = tmp
self.root = new_root
def double_left_rotate(self, new_root):
tmp_left = new_root.left
tmp_right = new_root.right
new_root.left = self.root.left
new_root.right = self.root
self.root.left.right = tmp_left
self.root.left = tmp_right
self.root = new_root
def single_right_rotate(self, new_root):
tmp = new_root.left
new_root.left = self.root
self.root.right = tmp
self.root = new_root
def double_right_rotate(self, new_root):
tmp_left = new_root.left
tmp_right = new_root.right
new_root.left = self.root
new_root.right = self.root.right
self.root.right.left = tmp_right
self.root.right = tmp_left
self.root = new_root
def insert(self, key):
node = self.root
direction = None
while node:
parent = node
if node.key == key:
print("Data already exists.")
return
elif node.key > key:
node = node.left
flag = "left"
if not direction:
direction = "left"
else:
node = node.right
flag = "right"
if not direction:
direction = "right"
new = Node(key)
if flag == "left":
parent.left = new
else:
parent.right = new
self.total += 1
self.get_max_depth(self.root)
if self.max_depth - self.balance() >= 1:
if direction == "left":
if self.root.left.key > key:
self.single_left_rotate(self.root.left)
else:
self.double_left_rotate(self.root.left.right)
else:
if self.root.right.key < key:
self.single_right_rotate(self.root.right)
else:
self.double_right_rotate(self.root.right.left)
def get_max_depth(self, tree, depth=0):
tmp = tree
if tmp == None:
depth -= 1
if self.max_depth < depth:
self.max_depth = depth
else:
self.get_max_depth(tmp.left, depth + 1)
self.get_max_depth(tmp.right, depth + 1)
def get_tree_structure(self, tree, depth):
tmp = tree
if tmp == None:
return
else:
self.get_tree_structure(tmp.left, depth + 1)
self.nodes[depth].append(tmp.key)
self.get_tree_structure(tmp.right, depth + 1)
def display_tree(self):
self.nodes = defaultdict(list)
self.get_tree_structure(self.root, 0)
depth = len(self.nodes)
for d in range(depth):
if d == 0:
print("root : ", end="")
else:
print("depth {}: ".format(d), end="")
for node in self.nodes[d]:
print("{0:3d}".format(node), end=", ")
print()
t = AVLTree(10)
print("insert 5, 20, 13, 30")
t.insert(5)
t.insert(20)
t.insert(13)
t.insert(30)
print("print tree...")
t.display_tree()
print("insert 15")
t.insert(15)
print("print tree...")
t.display_tree()
動作確認
それでは上記で作成したスクリプトを実行してみます。
> python avl_tree.py insert 5, 20, 13, 30 print tree... root : 10, depth 1: 5, 20, depth 2: 13, 30, insert 15 print tree... root : 13, depth 1: 10, 20, depth 2: 5, 15, 30,
今回は上記のfig1からfig2のような2重回転を想定してデータを挿入しました。
挿入前と挿入後に簡単に階層別に要素を表示して、バランスが取れていることが確認できます。
参考書籍
定本 Cプログラマのためのアルゴリズムとデータ構造 (SOFTBANK BOOKS)
