Rustで貪欲法 (3)

『アルゴリズムとデータ構造』を Rust で実装してみるシリーズ。7章の貪欲法。

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

今回は Multiple Array (AtCoder Grand Contest 009 A)

use proconio::input;

fn input(use_default_values: bool) -> (usize, Vec<u64>, Vec<u64>) {
    if use_default_values {
        let n = 7;
        let a = vec![3, 4, 5, 2, 5, 5, 9];
        let b = vec![1, 1, 9, 6, 3, 8, 7];
        return (n, a, b);
    } else {
        println!("Input n and a_b_pair:");
        input! {
            n: usize,
            a_b_pair: [(u64, u64); n],
        }
        let mut a: Vec<u64> = Vec::new();
        let mut b: Vec<u64> = Vec::new();
        for item in a_b_pair.into_iter() {
            a.push(item.0);
            b.push(item.1);
        }
        return (n, a, b);
    }
}

/// Multiple Array の解答例
pub fn code_7_3(use_default_values: bool) -> u64 {
    // 入力
    let (n, mut a, b) = input(use_default_values);

    // 答え
    let mut sum = 0;
    for i in (0..n).rev() {
        a[i] += sum; // 前回までの操作回数を足す
        let remainder = a[i] % b[i];
        let mut d = 0;
        if remainder != 0 {
            d = b[i] - remainder;
        }
        sum += d;
    }
    sum
}

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

    #[test]
    fn code_7_3_works() {
        assert_eq!(code_7_3(true), 22);
    }
}

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