NLL のおかげで Rust で平衡二分木を実装できた

Rust で平衡二分木を書くのは何となく難しいイメージがありました。 unsafe を使わずに実装できるものなのか気になったので、試しに実装してみました。

結論から言うと、unsafe を使わなくても平衡二分木は実装できました。また、unsafe だけでなく、RcRefCell も使っていません。 ただし NLL がないとコンパイルが通りませんでした。NLL は神機能ですね。

Splay

今回実装した平衡二分木は Splay 木というものです。

Splay 木は、search や insert などの典型的な操作を償却 O(\log n) 時間で実行できます。 また、Splay 木は最後にアクセスしたノードを木のルートに持ってくる性質があるため、アクセスに局所性がある場合によい性能が出るそうです。

平衡二分木といえば何と言っても赤黒木ですが、赤黒木は実装が大変すぎるので今回は却下しました。

アルゴリズム

Splay 木のアルゴリズムについては、Wikipedia に図付きで書いてあるので、こちらを参照下さい。

Splay tree - 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 は下図のように与えられた部分木のルートの左側のノードが新たなルートとなるように木を回転させる操作です。 もちろん、回転させるときにノードの順序関係は保つ必要があります。

f:id:nojima718:20181120012345p:plain

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 は最高。