・0 pv・3 min read
RustとAtCoderを勉強する(typical90_ca)
はじめに
AtCoder の問題を Rust で解いていきます。AtCoder も Rust も初心者ですが、温かい目で成長を見守っていただけるとありがたいです。
今回は、競プロ典型90問の079 - Two by Two（★3）を解きました。
提出コード
use itertools::{iproduct, Itertools};
use ndarray::{s, Array, Array2};
use proconio::input;
fn main() {
input! {
h: usize,
w: usize,
a: [[isize; w]; h],
b: [[isize; w]; h],
}
let a: Array2<isize> =
Array::from_shape_vec((h, w), a.into_iter().flatten().collect_vec()).unwrap();
let b: Array2<isize> =
Array::from_shape_vec((h, w), b.into_iter().flatten().collect_vec()).unwrap();
let mut diff = b - a;
let mut count = 0;
for (i, j) in iproduct!(0..h - 1, 0..w - 1) {
let d = diff[[i, j]];
diff.slice_mut(s![i..=i + 1, j..=j + 1])
.mapv_inplace(|x| x - d);
count += d.abs();
}
if diff.into_iter().all_equal() {
println!("Yes\n{}", count);
} else {
println!("No");
}
}解説
アルゴリズム的な解説はこちらを参考にしてください。この記事では Rust での実装に焦点を当てて解説します。今回は勉強がてら
ndarray を使用しました。Vec<Vec<_>>をArray2<_>に変換する
let a: Array2<usize> = Array::from_shape_vec((h, w), a.iter().flatten().collect()).unwrap();arrayの一部に対して演算を行う
diff.slice_mut(s![i..=i + 1, j..=j + 1]).mapv_inplace(|x| x - d);array の slice を取るメソッドはいくつかありますが、array の値を変えたいのであれば slice_mut を使います。さらに slice した範囲に演算を行う場合は map を続けますが、こちらも inplace とついたメソッドを使用しなければ array 自体の値は変わらないので注意してください。ちなみに map は参照を受け取り、 mapv は値を受け取ります。つまり、上のコードは下のように書き換えることもできます。diff.slice_mut(s![i..=i + 1, j..=j + 1]).map_inplace(|x| *x -= d);まとめ
Rust の
ndarray を使用してみました。さすがに numpy ほどの使いやすさではないと感じましたが、Vec よりは多次元配列の扱いがしやすくなると思うので、勉強していきたいと思います。