はじめに
はじめまして。名古屋大学3年のぱうえる（AtCoder: powell）です。最近はもっぱらRustで競プロをしています。
競プロとは
簡単に言えば、「与えられた問題を解くプログラムを作成したり、そのスピードを競ったりする競技」のことです。参加したことのない方はぜひ一度AtCoderのコンテストに出てみることをおすすめします。
さて、競プロの醍醐味といえばライブラリ整備ですが、Rust使いの競プロerが少ないせいか、Rustで書かれたライブラリはC++に比べてそう多くありません。また、Rustは所有権によってポインタの扱いが厳格に定まっているため、木のように再帰的なデータ構造では実装に少しコツがいります。
今回は、そんな困難を知ってなおRustでデータ構造を書きたいというRustaceanの一助となるべく、平衡二分探索木の中でも実装が簡単なスプレー木の実装を紹介していきたいと思います。
※自分もまだまだ勉強中の内容ですので、間違いを見つけた場合にはTwitter: @penguineer までご連絡いただけるとありがたいです。
説明
あまり厳密に書くと長くなってしまうので、ざっくり説明していきます。
二分探索木ってなあに？
二分探索木はデータ構造の一つで、主にプログラミング言語に組み込まれている集合型や辞書型などに用いられています。（C++言語のsetやmapなど。Pythonのdictはまた別のデータ構造）
配列などと比べたときの二分探索木の強みは、値を小さい順に並べておくことで、検索や挿入、削除を高速にできるという点です。
例
↓二分探索木では、予め要素を並べておくことで効率的に調べられる。
↓配列の場合、前から順番に調べていく必要がある。
スプレー木ってなあに？
二分探索木は、そのまま使っていると木の形に偏りが生まれて効率が悪くなる場合があります。これを回避するための機能がついているものが平衡二分探索木です。
平衡二分探索木も様々な種類があり、それぞれ色々な工夫をして木のバランスを保っています。
今回取り上げるスプレー木は、そのなかでもシンプルなアイデアで木の平衡を実現しているものです。
スプレー木の仕組み
シンプルとはいえ、ここで説明すると長くなるので省略します。ノードの動きについては、かつっぱさんの動画が非常にわかりやすいのでこちらをご参照ください。
実装の詳細はスプレー操作の実装を御覧ください。
実装
ノードの実装
今回はシンプルに、
- ノードが持つデータ
- 左の子へのポインタ
- 右の子へのポインタ
の3つのみを持つようにします。
pub struct SplayTreeNode<T> {
pub key: T, // ノードが持つデータ
left: Option<Box<SplayTreeNode<T>>>, // 左の子へのポインタ
right: Option<Box<SplayTreeNode<T>>>, // 右のへのポインタ
}なぜ片方向なのか
所有権のシステムによって、ループしたポインタを持つデータ構造はBoxのように単純なポインタ型ではなくRcとRefCellやunsafeな*mutなどで実装する必要があり少々複雑になってしまいます。簡単のため、今回は片方向で実装しました。
木全体の実装
根へのポインタと、木のサイズを持っておくようにします。
/// # SplayTree
/// スプレー木のクラス
pub struct SplayTree<T: Ord> {
size: usize, // 木のノード数を保存
pub root: Option<Box<Node<T>>>, // 根となるノードのポインタ
}回転の実装
親ノードを受け取り、回転したあとに親となるノードを返す形で実装します。
左回転
/// ## 左回転
fn rotate_left<T: Ord>(root: Option<Box<Node<T>>>) -> Option<Box<Node<T>>> {
if let Some(mut root) = root {
if let Some(mut new_root) = root.right {
root.right = new_root.left;
new_root.left = Some(root);
Some(new_root)
} else {
Some(root)
}
} else {
None
}
}右回転
/// ## 右回転
fn rotate_right<T: Ord>(root: Option<Box<Node<T>>>) -> Option<Box<Node<T>>> {
if let Some(mut root) = root {
if let Some(mut new_root) = root.left {
root.left = new_root.right;
new_root.right = Some(root);
Some(new_root)
} else {
Some(root)
}
} else {
None
}
}スプレー操作の実装
いよいよスプレー木の肝となるスプレー操作を実装していきます！
かつっぱさんの動画では、現在の頂点から親、親の親と上に見ていく実装でしたが、今回は子ノードへのポインタしか持っていないため、
- 現在の頂点
- 子
- 子の子
と下向きに三世代を見ていくようにします。（親は1つしかないのに対して、子は2ついることから場合分けは増えてしまいますが）
zig
目標となるノードが子ノードにあるときに行います。
- ターゲットが左の子：自分を右回転
- ターゲットが右の子：自分を左回転
zig-zig
目標となるノードが左の子の左の子、または右の子の右の子のときに行います。
- 左の子の左の子：自分を右回転→自分を右回転
- 右の子の右の子：自分を左回転→自分を左回転
zig-zag
目標となるノードが左の子の右の子、または右の子の左の子のときに行います
- 左の子の右の子：左の子を左回転→自分を右回転
- 右の子の左の子：右の子を右回転→自分を左回転
zig-zigステップと違って子の方を先に回転させるのに注意です！
スプレー全体
全体の実装はこんな感じになります。孫を再帰的にスプレーしておくのを忘れないようにしましょう。
実装
/// ## 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)
}
}ちゃんと使えるか確認
Rust playground
https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=4f60e6251c5bec982fa36763913b34a9AtCoder
https://atcoder.jp/contests/abc155/submissions/45743168Rust標準ライブラリのBTreeMapでは217msだったので、1.5倍弱。健闘してる方でしょうか。
おわりに
なかなかの重実装でしたが、いかがでしたか？噂によるとスプレー木が作れると他のつよつよライブラリも作りたい放題らしいので、みなさんもぜひチャレンジしてみてください！
次回のLinkCut木編でお会いしましょう！（未定です）
