Magicode logo
Magicode
・3 min read

RustとAtCoderを勉強する(typical90_ca)

AtCoderRust
https://cdn.apollon.ai/media/notebox/7863cc36-0805-4b2a-9e7b-cc46b21946d9.jpeg

はじめに

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 よりは多次元配列の扱いがしやすくなると思うので、勉強していきたいと思います。

Discussion

コメントにはログインが必要です。