NLL のおかげで Rust で平衡二分木を実装できた
Rust で平衡二分木を書くのは何となく難しいイメージがありました。 unsafe を使わずに実装できるものなのか気になったので、試しに実装してみました。
結論から言うと、unsafe を使わなくても平衡二分木は実装できました。また、unsafe だけでなく、Rc
や RefCell
も使っていません。
ただし NLL がないとコンパイルが通りませんでした。NLL は神機能ですね。
Splay 木
今回実装した平衡二分木は Splay 木というものです。
Splay 木は、search や insert などの典型的な操作を償却 時間で実行できます。 また、Splay 木は最後にアクセスしたノードを木のルートに持ってくる性質があるため、アクセスに局所性がある場合によい性能が出るそうです。
平衡二分木といえば何と言っても赤黒木ですが、赤黒木は実装が大変すぎるので今回は却下しました。
アルゴリズム
Splay 木のアルゴリズムについては、Wikipedia に図付きで書いてあるので、こちらを参照下さい。
重要なのは Splay 操作 という操作で、これは特定のキーを持つノードをルートに持ち上げていく操作です。 この持ち上げていく過程でその周辺のノードが何となく平衡状態に近づくようになっています。
search や insert は Splay 操作を使って簡単に実装できます。
実装
BinaryNode
ノードは key と左右のノードを保持します。今回は簡単のため value は保持しませんでした。 ノードは左右のノードを 所有 していることに注意してください。
struct BinaryNode<K: Ord> { key: K, left: Option<Box<BinaryNode<K>>>, right: Option<Box<BinaryNode<K>>>, }
SplayTree
SplayTree
は単に木のルートを保持しているだけです。
ノードが一つもない場合もあるので Option
になっています。
pub struct SplayTree<K: Ord> { root: Option<Box<BinaryNode<K>>>, }
rotate
データ構造の準備ができたので、Splay 操作の実装に入ります。
まずは rotate_right
です。
rotate_right
は下図のように与えられた部分木のルートの左側のノードが新たなルートとなるように木を回転させる操作です。
もちろん、回転させるときにノードの順序関係は保つ必要があります。
rotate_right
の実装は以下のようになりました。
// root の左側のノードが新たな根となるように木を回転させ、新たな根を返す。 fn rotate_right<K: Ord>(mut root: Box<BinaryNode<K>>) -> Box<BinaryNode<K>> { let mut new_root = root.left.unwrap(); root.left = new_root.right; new_root.right = Some(root); new_root }
まず、引数としてルートとなるノードの 値 を取ります。参照ではありません。
つまり、rotate_right
はルートノードの所有権を要求するということです。
そして、戻り値として新しいルートノードを所有権付きで返しています。
所有権が必要なのはノードの親子関係を弄るからです。 ノードの親子関係を変えるためには move する必要があり、move するにはそのノードを所有していないといけません。
また、関数の中身についてですが、実はこれ、NLL を有効にしないとコンパイルが通りません。 3行目でコンパイルエラーになります。
こういういかにも borrow check が難しそうなコードもちゃんとコンパイルしてくれる NLL は素晴らしいですね。
splay
rotate ができれば後は場合分けして splay 操作を実装するだけです。 key が左側にある場合と右側にある場合の2通り書かないといけないので、左側にある場合のみを示します。
// root を根とする部分木に対してスプレー操作を実行し、新たな根を返す。 // key を持つノードが部分木に存在する場合、それが新たな根となる。 // key が部分木に存在しない場合、最後にアクセスしたノードが根となる。 // 部分木は破壊的に変更される。 fn splay<K: Ord>(key: &K, root: Option<Box<BinaryNode<K>>>) -> Option<Box<BinaryNode<K>>> { if root.is_none() { return None; } let root = root.unwrap(); let new_root = match key.cmp(&root.key) { Ordering::Less => splay_left(key, root), Ordering::Greater => splay_right(key, root), Ordering::Equal => root, }; Some(new_root) } // key が root の左側にあるときのスプレー操作を行う。新たな根を返す。 fn splay_left<K: Ord>(key: &K, mut root: Box<BinaryNode<K>>) -> Box<BinaryNode<K>> { if root.left.is_none() { return root; } let mut left = root.left.unwrap(); if key < &left.key { // zig-zig // left-left の部分木の根に key を持ってくる left.left = splay(key, left.left); root.left = Some(left); // 右回転を2回行って left-left を根に持ってくる let new_root = rotate_right(root); if new_root.left.is_some() { rotate_right(new_root) } else { new_root } } else if key > &left.key { // zig-zag // left-right の部分木の根に key を持ってくる left.right = splay(key, left.right); // 左回転と右回転を行って left-right を根に持ってくる root.left = if left.right.is_some() { Some(rotate_left(left)) } else { Some(left) }; rotate_right(root) } else { // zig root.left = Some(left); rotate_right(root) } }
insert とか search とか
長くなってきたので、ここでは省略します。GitHub に全体のソースコードを上げているので、そっちを参照してください。
workspace/splay_tree.rs at master · nojima/workspace · GitHub
まとめ
NLL は最高。