『アルゴリズムとデータ構造』を Rust で実装してみるシリーズ。6章の二分探索法。
前回 → https://quantized-cube.com/rust-binary-search-03
今回は射撃王 (AtCoder Beginner Contest 023 D)
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
コメント