Rustで二分探索法 (4)

Rust

2021.12.27 2021.11.16

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

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

use proconio::input;
use std::cmp;

fn input(use_default_values: bool) -> (usize, Vec<usize>, Vec<usize>) {
    if use_default_values {
        let n = 4;
        let h = vec![5, 12, 14, 21];
        let s = vec![6, 4, 7, 2];
        return (n, h, s);
    } else {
        println!("Input n and h_s_pair:");
        input! {
            n: usize,
            h_s_pair: [(usize, usize); n],
        }
        let mut h: Vec<usize> = Vec::new();
        let mut s: Vec<usize> = Vec::new();
        for item in h_s_pair.into_iter() {
            h.push(item.0);
            s.push(item.1);
        }
        return (n, h, s);
    }
}

/// 射撃王問題に対する二分探索法
pub fn code_6_5(use_default_values: bool) -> usize {
    // 入力
    let (n, h, s) = input(use_default_values);

    // 二分探索の上限値を求める
    let mut m: usize = 0;
    for i in 0..n as usize {
        m = cmp::max(m, h[i] + s[i] * n);
    }

    // 二分探索
    let mut left = 0;
    let mut right = m;

    while right - left > 1 {
        let mid = left + (right - left) / 2;
        // 判定する
        let mut ok = true;
        let mut t = vec![0; n]; // 各風船を割るまでの制限時間
        for i in 0..n as usize {
            // そもそも mid が初期高度より低かったら false
            if mid < h[i] {
                ok = false;
            } else {
                t[i] = (mid - h[i]) / s[i];
            }
        }

        // 時間制限がさし迫っている順にソート する
        t.sort();
        for i in 0..n {
            // 時間切れ発生の場合は false
            if t[i] < i {
                ok = false;
            }
        }

        if ok {
            right = mid;
        } else {
            left = mid;
        }
    }

    right
}

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

    #[test]
    fn code_6_5_works() {
        assert_eq!(code_6_5(true), 23);
    }
}

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