Rust 関数などの引数にRangeを受け取る

(2022/07/17)用途別の記事を作成しました。

おそらく上の記事の方がおすすめです。

Range 構造体について

Rust の Range

Range構造体は、for文などでよく使っている0..10のようなものです。
この構造体は、std::ops::Rangeです(ドキュメント)。

for i in 0..10 {
    println!("{}", i);
}

Range を引数として受け取るときにどう書くか

関数を書いていると、Rangeを引数として受け取ることができるときれいに書くことができるのに、と思って検索してみてもいまいち情報を見つけることができない、という経験があるのは私だけではないと思います。

std::vec::Vecなどにもこの Range を引数として受け取ることができるようなメソッドがいくつかあるので参考にできるのですが(たとえばdrainメソッド)、std::ops::RangeBoundsトレイトを使って、次のように書くことができます。

use std::ops::{Bound, RangeBounds};
pub fn range_decide<R: RangeBounds<usize>>(range: R, len: usize) -> (usize, usize) {
    let start = match range.start_bound() {
        Bound::Unbounded => 0,
        Bound::Excluded(&s) => s + 1,
        Bound::Included(&s) => s,
    };
    let end = match range.end_bound() {
        Bound::Unbounded => len,
        Bound::Excluded(&t) => t.min(len),
        Bound::Included(&t) => (t + 1).min(len),
    };
    assert!(start <= end);
    (start, end)
}

この関数はRangeと、長さであるlenを受け取り、その範囲の両端を返すようなものです。文字で説明するよりも実際の動作を見た方が分かりやすいと思うので、実行例をいくつか載せます。(Rust Playground)

fn main() {
    let (start, end) = range_decide(0..10, 100);
    println!("start = {}, end = {}", start, end);  // start = 0, end = 10
    let (start, end) = range_decide(0.., 100);
    println!("start = {}, end = {}", start, end);  // start = 0, end = 100
    let (start, end) = range_decide(5..=10, 100);
    println!("start = {}, end = {}", start, end);  // start = 5, end = 11
    let (start, end) = range_decide(..10, 100);
    println!("start = {}, end = {}", start, end);  // start = 0, end = 10
    let (start, end) = range_decide(.., 100);
    println!("start = {}, end = {}", start, end);  // start = 0, end = 100
    let (start, end) = range_decide(10..10, 100);
    println!("start = {}, end = {}", start, end);  // start = 10, end = 10
}

同様の処理はこのようにmatch式で頑張らなくてもassert_lenメソッド(ドキュメント)で実現できるのですが、残念ながらこのメソッドは現時点ではまだ安定版では使用できず、実験的なものだそうです(2021/02現在 Rust 1.50.0)。

使用例

二分探索

同様な処理を用いて、二分探索のコードを書いてみます。二分探索と言ってもソート済み配列から値を探すものではなく、もう少し一般的なもので、ある整数 $X$ について、$x \le X-1$ のときに $f(x)=false$ 、$x \ge X$ のときに $f(x)=true$ となるような関数 $f(x)$ について、境目である $X$ を求めるようなものです。

\begin{eqnarray}
f( x ) = \begin{cases}
        false & ( x \le X-1 ) \\
        true & ( X \le x )
    \end{cases}
\end{eqnarray}

一般的な二分探索のより詳しい説明は、他の方がいろいろな記事を書いてくださっています(参考: 二分探索アルゴリズムを一般化 〜 めぐる式二分探索法のススメ 〜)。
さて、上の実装は簡単のためusizeの範囲しか取れないような書き方でしたが、今回はせっかくなので、numクレートを使ってもう少し抽象的に書きたいと思います。

実際のコード

use std::ops::{AddAssign, Bound, RangeBounds};

use num::Num;

/// **O(log(end-start))**, return (start, end) that f(end) == true (should f(start) == false)
pub fn initial_indices<I, R, F>(range: R, f: &F) -> (I, I)
where
    I: Clone + PartialOrd + AddAssign + Num,
    R: RangeBounds<I>,
    F: Fn(&I) -> bool,
{
    let start = match range.start_bound() {
        Bound::Unbounded => I::zero(),
        Bound::Excluded(start) => (start.clone() + I::one()),
        Bound::Included(start) => start.clone(),
    };
    let (start, end) = match range.end_bound() {
        Bound::Unbounded => {
            let (mut guessed_start, mut range) = (start.clone(), I::one());
            let guessed_end = loop {
                let end = start.clone() + range.clone();
                if f(&end) {
                    break end;
                } else {
                    guessed_start = end;
                    range += range.clone();
                }
            };
            (guessed_start, guessed_end)
        }
        Bound::Excluded(end) => {
            (start, if end > &I::one() { end.clone() - I::one() } else { I::zero() })
        }
        Bound::Included(end) => (start, end.clone()),
    };
    (start, end)
}

/// **O(log(ans))**, find the first index at which false -> true (f(start) must be false)
pub fn bisect<I, R, F>(range: R, f: F) -> Option<I>
where
    I: Clone + PartialOrd + AddAssign + Num,
    R: RangeBounds<I>,
    F: Fn(&I) -> bool,
{
    let (mut start, mut end) = initial_indices(range, &f);
    if start >= end || f(&start) || !f(&end) {
        return None; // if f(start) == true then all is true, and if f(end)==false then all is false
    }
    while start.clone() + I::one() < end {
        let mid = (start.clone() + end.clone()) / (I::one() + I::one());
        if f(&mid) {
            end = mid;
        } else {
            start = mid;
        }
    }
    Some(end)
}

すこしごちゃごちゃしてしまったので、簡単に説明します。上のinitial_indices関数が引数で与えられたRange構造体をもとに、探索範囲を決定する関数です。上限が決められていなかった場合は、 $f(end)$ がtrueになるまで、探索範囲を倍、倍と広げていきます。下のbisect関数では二分探索によって $f(x)$ がfalseで $f(x+1)$ がtrueとなるような境目を探ります。このとき、startは $f(start)$ が常にfalseになるように、endは $f(end)$ が常にtrueになるように、範囲を絞っていくことで、最終的なendが $f(x)$ がtrueになる初めの $x$ ということになります。

使ってみる

それでは、上のbisect関数を用いて、 $y = x^2$ について、$x \ge 0$ の範囲で、初めて $y \ge 100$ となるような整数 $x$ を求めてみます(答えは当然 $x = 10$ ですが (Rust Playground))。

fn main() {
    let x_pow_2 = |x| x * x;
    println!("{:?}", bisect(0.., |&i| x_pow_2(i) >= 100))  // Some(10)
}

関数の呼び出し側を見ると、Range構造体を使って、「$x \ge 0$ の範囲」をきれいに表現できています。

まとめ

RangeBoundsを使うことで、引数にRangeをとるような関数をいい感じに書くことができます。
応用するとダイクストラの引数のstart、endもRangeを使うことでいい感じに表現できるかもしれないですね(endを指定しない場合はstartから連結成分すべてについて最短距離を求めるようにするなど(これに関してはRangeに求めることを少し逸脱しているかもしれないです))。