Search code examples
matrixrustvectorrangemapping

Map range to array in rust (nested)

The following rust code calculates the product matrix * vector as an array on the stack and displays it:

fn main () {
    let matrix: [[f64;3];3] = [
        [ 1.0, 0.0, 0.0 ],
        [ -1.0, 1.0, 0.0 ],
        [ 1.0, 0.0, -1.0 ],
    ];
    let vector: [f64;3] = [
        3.0,
        -2.0,
        1.0,
    ];
    let range = [ 0, 1, 2 ];
    let result: [f64;3] = range.map(
        |i| range.map(
            |j| matrix[i][j] * vector[j]
        ).iter().sum()
    );
    println!( "{:?}", result );
}

In principle this is fine, but I find the use of an array of indices awkward. How could I use a range instead for the indices like this:

    let range = 0..3;
Solution

  • To construct an array elementwise using a function, use from_fn:

    let result: [f64; 3] = std::array::from_fn(
    |i| (0..3).map(
        |j| matrix[i][j] * vector[j]
    ).sum()
);

    Playground