Rustでデータ構造書いてみた（スプレー木編）

簡単に言えば、「与えられた問題を解くプログラムを作成したり、そのスピードを競ったりする競技」のことです。参加したことのない方はぜひ一度AtCoderのコンテストに出てみることをおすすめします。

競プロに関しての記事は前々日と前日にも公開されているのでぜひご覧ください。

所有権のシステムによって、ループしたポインタを持つデータ構造はBoxのように単純なポインタ型ではなくRcとRefCellやunsafeな*mutなどで実装する必要があり少々複雑になってしまいます。簡単のため、今回は片方向で実装しました。

/// ## splay_inner
/// splay操作を行う
/// ### 戻り値
/// - `Option<Box<Node<T>>>`：新しく根となるノード
/// - `bool`：目的の値が存在したかどうか
fn splay_inner<T: Ord>(
    mut root: Option<Box<Node<T>>>,
    key: &T,
) -> (Option<Box<Node<T>>>, bool) {
    if root.is_none() {
        return (root, false);
    }
    // 孫 → 子
    match key.cmp(&root.as_deref().unwrap().key) {
        Ordering::Equal => (root, true),
        Ordering::Less => {
            // 左の子
            let left = &mut root.as_deref_mut().unwrap().left;
            if left.is_none() {
                return (root, false);
            }
            match key.cmp(&left.as_deref().unwrap().key) {
                // zig
                Ordering::Equal => {
                    // 左の子をrootに
                    (rotate_right(root), true)
                }
                // zig-zig
                Ordering::Less => {
                    // 孫をsplay
                    let left_left = left.as_deref_mut().unwrap().left.take();
                    let (mut new_left_left, is_found) = splay_inner(left_left, key);
                    swap(&mut left.as_deref_mut().unwrap().left, &mut new_left_left);
                    // 親を右に回転
                    let tmp_child = rotate_right(root);
                    // さらに右に回転
                    (rotate_right(tmp_child), is_found)
                }
                // zig-zag
                Ordering::Greater => {
                    // 孫をsplay
                    let left_right = left.as_deref_mut().unwrap().right.take();
                    let (mut new_left_right, is_found) = splay_inner(left_right, key);
                    swap(&mut left.as_deref_mut().unwrap().right, &mut new_left_right);
                    // 左の子を左に回転
                    let left = root.as_deref_mut().unwrap().left.take();
                    let mut new_left = rotate_left(left);
                    swap(&mut root.as_deref_mut().unwrap().left, &mut new_left);
                    // さらに右に回転
                    (rotate_right(root), is_found)
                }
            }
        }
        Ordering::Greater => {
            // 右の子
            let right = &mut root.as_deref_mut().unwrap().right;
            if right.is_none() {
                return (root, false);
            }
            match key.cmp(&right.as_deref().unwrap().key) {
                // zig
                Ordering::Equal => {
                    // 右の子をrootに
                    (rotate_left(root), true)
                }
                // zig-zag
                Ordering::Less => {
                    // 孫をsplay
                    let right_left = right.as_deref_mut().unwrap().left.take();
                    let (mut new_right_left, is_found) = splay_inner(right_left, key);
                    swap(&mut right.as_deref_mut().unwrap().left, &mut new_right_left);
                    // 右の子を右に回転
                    let right = root.as_deref_mut().unwrap().right.take();
                    let mut new_right = rotate_right(right);
                    swap(&mut root.as_deref_mut().unwrap().right, &mut new_right);
                    // さらに左に回転
                    (rotate_left(root), is_found)
                }
                // zig-zig
                Ordering::Greater => {
                    // 孫をsplay
                    let right_right = right.as_deref_mut().unwrap().right.take();
                    let (mut new_right_right, is_found) = splay_inner(right_right, key);
                    swap(
                        &mut right.as_deref_mut().unwrap().right,
                        &mut new_right_right,
                    );
                    // 親を左に回転
                    let tmp_child = rotate_left(root);
                    // さらに左に回転
                    (rotate_left(tmp_child), is_found)
                }
            }
        }
    }
}

impl SplayTree<T> {
		/// ## splay
    /// スプレー操作をおこなう
    /// ### 戻り値
    /// - `bool`：要素が存在したかどうか
    pub fn splay(&mut self, key: &T) -> bool {
        // 根の取り出し
        let root = self.root.take();
        // スプレー操作
        let (new_root, is_found) = splay_inner(root, key);
        self.root = new_root;
        is_found
    }

		/// ## get
		/// 値の検索を行う
		/// ### 戻り値
		/// - `Option<&T>`: キーに紐づいた値
		pub fn get(&mut self, key: &T) -> Option<&T> {
		    if self.splay(key) {
		        Some(&self.root.as_deref().unwrap().key)
		    } else {
		        None
		    }
		}
}

impl SplayTree<T> {
		/// ## insert
    /// 値の挿入を行う。
    /// すでに同じキーが存在した場合は値を置き換えて前の値を返す。
    /// ### 戻り値
    /// - `Option<T>`: 以前の値
    pub fn insert(&mut self, key: T) -> Option<T> {
        // rootの取り出し
        let root = self.root.take();
        // splay操作
        let (mut tmp_root, is_found) = splay_inner(root, &key);
        if is_found {
            self.root = tmp_root;
            let res = replace(&mut self.root.as_deref_mut().unwrap().key, key);
            return Some(res);
        }
        // 挿入
        self.root = Some(Box::new(Node::new(key.clone())));
        if tmp_root.is_some() {
            match key.cmp(&tmp_root.as_deref().unwrap().key) {
                Ordering::Equal => unreachable!(),
                Ordering::Less => {
                    let mut new_left = tmp_root.as_deref_mut().unwrap().left.take();
                    swap(&mut self.root.as_deref_mut().unwrap().left, &mut new_left);
                    swap(&mut self.root.as_deref_mut().unwrap().right, &mut tmp_root);
                }
                Ordering::Greater => {
                    let mut new_right = tmp_root.as_deref_mut().unwrap().right.take();
                    swap(&mut self.root.as_deref_mut().unwrap().right, &mut new_right);
                    swap(&mut self.root.as_deref_mut().unwrap().left, &mut tmp_root);
                }
            }
        }
        // 要素数の更新
        self.size += 1;
        None
    }
}

impl SplayTree<T> {
		/// ## delete
    /// 値の削除
    /// ### 戻り値
    /// - `Option<T>`: 削除された値
    pub fn delete(&mut self, key: &T) -> Option<T> {
        // rootの取り出し
        let root = self.root.take();
        // splay操作
        let (mut tmp_root, is_found) = splay_inner(root, key);
        if !is_found {
            self.root = tmp_root;
            return None;
        }
        // 削除
        if tmp_root.as_deref().unwrap().left.is_none() {
            swap(&mut self.root, &mut tmp_root.as_deref_mut().unwrap().right);
        } else {
            let root_left = tmp_root.as_deref_mut().unwrap().left.take();
            swap(&mut self.root, &mut splay_inner(root_left, key).0);
            swap(
                &mut self.root.as_deref_mut().unwrap().right,
                &mut tmp_root.as_deref_mut().unwrap().right,
            );
        }
        let deleted = tmp_root.take();
        // 要素数の更新
        self.size -= 1;
        Some(deleted.unwrap().key)
    }
}