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

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

https://qiita.com/hystcs/items/8e064f8b9a79adb9cca7

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

# Range 構造体について
## Rust の Range
Range構造体は、for文などでよく使っている`0..10`のようなものです。
この構造体は、`std::ops::Range`です([ドキュメント](https://doc.rust-lang.org/std/ops/struct.Range.html))。

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

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

`std::vec::Vec`などにもこの Range を引数として受け取ることができるようなメソッドがいくつかあるので参考にできるのですが([たとえば`drain`メソッド](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.drain))、`std::ops::RangeBounds`トレイトを使って、次のように書くことができます。

```Rust
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](https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=beb648170fc846c63d6857f0976afe7e))

```Rust
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`メソッド([ドキュメント](https://doc.rust-lang.org/std/ops/trait.RangeBounds.html#method.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$ を求めるようなものです。

```math
\begin{eqnarray}
f( x ) = \begin{cases}
        false & ( x \le X-1 ) \\
        true & ( X \le x )
    \end{cases}
\end{eqnarray}
```
一般的な二分探索のより詳しい説明は、他の方がいろいろな記事を書いてくださっています(参考: [二分探索アルゴリズムを一般化 〜 めぐる式二分探索法のススメ 〜](https://qiita.com/drken/items/97e37dd6143e33a64c8c))。
さて、上の実装は簡単のため`usize`の範囲しか取れないような書き方でしたが、今回はせっかくなので、numクレートを使ってもう少し抽象的に書きたいと思います。

### 実際のコード
```Rust
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](https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=922a21596fc4710ada75e3b010edd669)))。

```Rust
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に求めることを少し逸脱しているかもしれないです))。