I have two 2D arrays in Rust like this

    let d1: [[T; 3]; 3] = ...
    let d2: [[T; 3]; 3] = ...

Now I want to add the elements element-wise and assign the result to a new 2D array.

My current solution looks like this:

    let mut result = [[T::default(); ROW]; COL];
    
    for (i, row) in d1.iter().enumerate() {
        for (j, element) in row.iter().enumerate() {
            result[i][j] = *element   d2[i][j];
        }
    }

This is inefficient as it first fills result with default values only to override them later.

Therefore I want to solve this using iterators by flattening the 2D arrays and then adding the elements.

However, I don't know how I can construct a 2D array from the resulting flattened 1D iterator :

    let result = for sum in d1.iter().flatten().zip(
                            d2.iter().flatten())
                            .map(|(x, y)| x y)
                            .collect().???

CodePudding user response：

This code adds a 2D array without initializing it with default elements, but it does go through a Vec to array conversion, because you cannot collect directly into an array. I estimate the likelihood of this being optimized away to be much lower than for your default elements solution.

fn add_3by3(a: [[u32; 3]; 3], b: [[u32; 3]; 3]) -> [[u32; 3]; 3] {
    a.iter()
        .zip(b.iter())
        .map(|(ai, bi)| {
            ai.iter()
                .zip(bi.iter())
                .map(|(aij, bij)| aij   bij)
                .collect::<Vec<_>>()
                .try_into()
                .unwrap()
        })
        .collect::<Vec<_>>()
        .try_into()
        .unwrap()
}

Incidentally, I think it is probably a better idea to keep your data in a 1D array no matter what actual shape it has, and keep shape information on the side. That way every element-wise operation is very simple.

CodePudding user response：

You can do a hybrid approach.

You can initialize a fixed-size array via first constructing a [(); N] and then calling .map() on it. Combined with iterators over the input arrays, this should be somewhat performant.

Although I still recommend benchmarking it vs the naive solution.

fn add_2d<'a, 'b, T, const W: usize, const H: usize>(
    d1: &'a [[T; W]; H],
    d2: &'b [[T; W]; H],
) -> [[T; W]; H]
where
    T: 'static,
    &'a T: std::ops::Add<&'b T, Output = T>,
{
    let mut iter_row = d1.iter().zip(d2.iter());

    [(); H].map(move |()| {
        let (row1, row2) = iter_row.next().unwrap();
        let mut elem_iter = row1.iter().zip(row2.iter());
        [(); W].map(move |()| {
            let (elem1, elem2) = elem_iter.next().unwrap();
            elem1   elem2
        })
    })
}

fn main() {
    let d1: [[u32; 2]; 3] = [[1, 2], [4, 5], [7, 8]];
    let d2: [[u32; 2]; 3] = [[10, 20], [40, 50], [70, 80]];

    let sums = add_2d(&d1, &d2);
    println!("{:?}", sums);
}

[[11, 22], [44, 55], [77, 88]]

Alternatively, if you don't mind unsafe, this is what I would probably go for to achieve maximum performance:

use std::mem::MaybeUninit;

fn add_2d<'a, 'b, T, const W: usize, const H: usize>(
    d1: &'a [[T; W]; H],
    d2: &'b [[T; W]; H],
) -> [[T; W]; H]
where
    T: 'static,
    &'a T: std::ops::Add<&'b T, Output = T>,
{
    let mut result: MaybeUninit<[[T; W]; H]> = MaybeUninit::uninit();

    let arr_ptr = result.as_mut_ptr() as *mut [T; W];

    unsafe {
        for y in 0..H {
            let row_ptr = arr_ptr.add(y) as *mut T;
            for x in 0..W {
                row_ptr.add(x).write(&d1[y][x]   &d2[y][x]);
            }
        }

        result.assume_init()
    }
}

fn main() {
    let d1: [[u32; 2]; 3] = [[1, 2], [4, 5], [7, 8]];
    let d2: [[u32; 2]; 3] = [[10, 20], [40, 50], [70, 80]];

    let sums = add_2d(&d1, &d2);
    println!("{:?}", sums);
}

[[11, 22], [44, 55], [77, 88]]

Be sure to use .write() instead of a simple assignment operation, because the assignment operation would run the Drop function of the previous, uninitialized element. Which would cause undefined behaviour.