Rustで二分探索法 (1)

『アルゴリズムとデータ構造』を Rust で実装してみるシリーズ。6章の二分探索法。

前回 → https://quantized-cube.com/binary-search-rust

今回は二分探索法と「一般化した二分探索法」を実装する。後者はハードコーディングしてるけど、そもそも元のコードもフレームを示しているだけなので。。

use proconio::input;


fn input_vector(use_default_values: bool) -> (usize, Vec<i64>) {
    if use_default_values {
        let n = 8;
        let a = vec![3, 5, 8, 10, 14, 17, 21, 39];
        return (n, a);
    } else {
        println!("Input n and a:");
        input! {
            n: usize,
            a: [i64;n],
        }
        return (n, a);
    }
}

fn input_key() -> i64 {
    println!("Input key:");
    input! {
        key: i64,
    }
    key
}

/// 目的の値 key の添字を返す (存在しない場合は -1)
fn binary_search_6_1(a: &Vec<i64>, key: i64) -> i64 {
    // 配列 a の左端と右端
    let mut left = 0 as i64;
    let mut right = (a.len() as i64) - 1;

    while left <= right {
        let mid = left + (right - left) / 2; // 区間の真ん中
        let a_mid = a[mid as usize];
        if a_mid == key {
            return mid;
        } else if a_mid > key {
            right = mid - 1;
        } else if a_mid < key {
            left = mid + 1;
        }
    }
    -1
}

/// 配列から目的の値を探索する二分探索法
/// key が None なら標準入力から入力する
pub fn code_6_1(key: Option<i64>, use_default_values: bool) -> i64 {
    let (_, a) = input_vector(use_default_values);
    let key = key.unwrap_or_else(|| input_key());
    binary_search_6_1(&a, key)
}

/// x が条件を満たすかどうか
fn p(x: i64) -> bool {
    let (_, a) = input_vector(true);
    let key = 9;
    if a[x as usize] >= key {
        true
    } else {
        false
    }
}

/// P(x) = true となる最小の整数 x を返す
fn binary_search_6_2() -> i64 {
    // P(left) = False, P(right) = True となるように
    let mut left = 0 as i64;
    let mut right = 7 as i64;
    while right - left > 1 {
        let mid = left + (right - left) / 2;
        if p(mid) {
            right = mid;
        } else {
            left = mid;
        }
    }
    right
}

/// 一般化した二分探索法の基本形
pub fn code_6_2() -> i64 {
    binary_search_6_2()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn code_6_1_works() {
        assert_eq!(code_6_1(Some(10), true), 3);
        assert_eq!(code_6_1(Some(3), true), 0);
        assert_eq!(code_6_1(Some(39), true), 7);

        assert_eq!(code_6_1(Some(-100), true), -1);
        assert_eq!(code_6_1(Some(9), true), -1);
        assert_eq!(code_6_1(Some(100), true), -1);
    }

    #[test]
    fn code_6_2_works() {
        assert_eq!(code_6_2(), 3);
    }
}

GitHub のリポジトリはこちら https://github.com/quantized-cube/book_algorithm_solution_rust