【rust】AtCoderで使う二次元ndarrayの作り方

モチベーション

使えれば便利そうなndarrayクレートですが、二次元のVecをndarrayにする方法が調べてもよくわからなかったので、記事にしました。

AtCoderをやっていると入力の受け取りにproconioを使うと思います。proconioを使って二次元のVec(Vec<Vec<T>>)をよく受け取るので、これを二次元のndarrayにしたいのです。

ndarray::ArrayBase - Rust

API documentation for the Rust `ArrayBase` struct in crate `ndarray`.

docs.rs

Vec<Vec<T>>からndarrayへ

結論から書くと、以下のように変換します。

ここで、aはVec<Vec<T>>でr行c列です。

let a = Array::from_shape_vec((r,c), a.concat()).unwrap();

以下、それぞれのメソッドの説明を書いていきます。

Array::from_shape_vec

このメソッドは、Vec<T>を任意の形のArrayを作ってくれます。1つ目の引数に形を、2つ目の引数に一次元のVecを指定します。ここでは二次元にしたいので、(r,c)としています。

concat

concatを使うと、Vecを一次元にすることができます。

これはstdに入っています。

Vec in std::vec - Rust

A contiguous growable array type, written as `Vec`, short for ‘vector’.

doc.rust-lang.org

例

use ndarray::prelude::*;

fn main() {
    let a = vec![vec![0, 1, 2], vec![3, 4, 5], vec![6, 7, 8]];
    let a = Array::from_shape_vec((3,3),  a.concat()).unwrap();
    eprintln!("a = {:?}", a);
}

実行結果

a = [[0, 1, 2],
     [3, 4, 5],
     [6, 7, 8]], shape=[3, 3], strides=[3, 1], layout=Cc (0x5), const ndim=2