Rustで二分探索法 (3)

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

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

今回はペア和を最適化する問題。

use proconio::input;

const INF: i64 = 20000000; // 十分大きな値に

fn input(use_default_values: bool) -> (usize, i64, Vec<i64>, Vec<i64>) {
    if use_default_values {
        let n = 3;
        let k = 10;
        let a = vec![8, 5, 4];
        let b = vec![4, 1, 9];
        return (n, k, a, b);
    } else {
        println!("Input n, k, a, b:");
        input! {
            n: usize,
            k: i64,
            a: [i64; n],
            b: [i64; n],
        }
        return (n, k, a, b);
    }
}

/// P(x) = true となる最小の整数 x を返す
fn lower_bound<T>(v: &Vec<T>, key: T) -> T
where
    T: Ord + Copy,
{
    let mut left = 0 as usize;
    let mut right = v.len();
    while right - left > 1 {
        let mid = left + (right - left) / 2;
        if v[mid] >= key {
            right = mid;
        } else {
            left = mid;
        }
    }
    v[right]
}

/// 二分探索法を用いて、「ペア和を最適化する問題」
/// に対する全探索解法を高速化する
pub fn code_6_4(use_default_values: bool) -> i64 {
    // 入力を受け取る
    let (n, k, a, b) = input(use_default_values);
    let mut b = b;

    // 暫定最小値を格納する変数
    let mut min_value = INF;

    // b をソート
    b.sort();

    // b に無限大を表す値 (INF) を追加しておく
    // これを行うことで、val = b.last() となる可能性を除外する
    b.push(INF);

    // a を固定して解く
    for i in 0..n {
        // b の中で k - a[i] 以上の範囲での最小値
        let val = lower_bound(&b, k - a[i]);

        // min_value と比較する
        if a[i] + val < min_value {
            min_value = a[i] + val;
        }
    }

    min_value
}

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

    #[test]
    fn code_6_4_works() {
        assert_eq!(code_6_4(true), 12);
    }
}

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