Ccmmutty logo
Commutty IT
・2 min read

RustとAtCoderを勉強する(typical90_d)

AtCoderRust
https://cdn.magicode.io/media/notebox/ebf0e01a-dbcd-4127-a7b8-03063e027b78.jpeg

はじめに

AtCoder の問題を Rust で解いていきます。AtCoder も Rust も初心者ですが、温かい目で成長を見守っていただけるとありがたいです。
今回は、競プロ典型90問の004 - Cross Sum（★2）を解きました。

提出コード

rust
use itertools::Itertools;
use ndarray::{Array, Array2};
use proconio::fastout;

#[fastout]
fn main() {
    proconio::input! {
        h: usize,
        w: usize,
        a: [[usize; w]; h],
    }
    let a: Array2<usize> =
        Array::from_shape_vec((h, w), a.into_iter().flatten().collect()).unwrap();

    let sum_row = (0..h).map(|i| a.row(i).sum()).collect_vec();
    let sum_col = (0..w).map(|j| a.column(j).sum()).collect_vec();

    for i in 0..h {
        let b = (0..w)
            .map(|j| sum_row[i] + sum_col[j] - a[[i, j]])
            .join(" ");
        println!("{}", b);
    }
}

解説

H行W列の行列が与えられ、その行列の全てのマスに対して 同じ行または同じ列（自分自身を含む） の値の合計を求める問題です。自分と同じ行と列の合計を足し合わせ、自分自身の値を引けば良いです。ただし、全てのマスで行・列の合計を計算すると時間が足りなくなることに注意です。
また、2次元の配列を扱うということで今回は勉強がてら ndarray を使用しました。proconio::input! では ndarray::Array で読み込めなかったため、Vec<Vec<usize>> で入力を読み込んだ後に変換しました。ndarray::Array で直接読み込む方法はあるんでしょうか？有識者の方教えてください...。

まとめ

  • 2次元 Vec の1次元化は vec.into_iter().flatten().collect() とする
  • メソッドチェーン楽しい

Discussion

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