CODEWITHSUNDEEP
arraysrust
user1569030
user1569030

Reputation:

How do I add two Rust arrays element-wise?

This is an absolute beginner question, but I can't find anything useful after searching for half an hour. I have Rust 1.7.0 and this code:

type coord = [i64; 3];

// add two coordinates ("vectors") pointwise, that is
// if z = add(a, b) then z[i] = a[i] + b[i] for i=0..2
fn add(a: coord, b: coord) -> coord {
    //???
}

The obvious thing I tried first is a.zip(b).map(|(u,v)| u+v) but this doesn't work (can't zip arrays), nor does a.iter().zip(b.iter()).map(|(u,v)| u+v) because it can't convert the iterator back to an array. I can see why this doesn't work in general but in this case we know both things are the same length.

For now I'm doing

fn add(a: coord, b: coord) -> coord {
    let mut z: coord = [0, 0, 0];
    for i in 0..2 {
        z[i] = a[i] + b[i];
    }
    z
}

but it looks ugly by comparison. What am I missing?

Upvotes: 9

Views: 15629

Answers (1)

user4815162342
user4815162342

Reputation: 155600

One simple approach is to generate indices using the enumerate iterator method and fill z the "obvious" way, by assigning into the obtained indices:

type Coord = [i64; 3];

fn add(a: Coord, b: Coord) -> Coord {
    let mut z: Coord = [0, 0, 0];
    for (i, (aval, bval)) in a.iter().zip(&b).enumerate() {
        z[i] = aval + bval;
    }
    z
}

fn main() {
    let x: Coord = [1, 2, 3];
    let y: Coord = [1, 1, 1];
    assert!(add(x, y) == [2, 3, 4]);
}

In Rust, we can do better than that by noticing that iter() produces references into the array, and iter_mut() is available to produce mutable references. This results in code very similar to what you attempted to write:

fn add(a: Coord, b: Coord) -> Coord {
    let mut z: Coord = [0, 0, 0];
    for ((zref, aval), bval) in z.iter_mut().zip(&a).zip(&b) {
        *zval = aval + bval;
    }
    z
}

If this pattern of writing into z recurs with different operations, you can abstract the creation of new Coord and filling it with data into a generic function:

fn new_coord_from<F: Iterator<Item=i64>>(src: F) -> Coord {
    let mut result = [0; 3];
    for (rref, val) in result.iter_mut().zip(src) {
        *rref = val;
    }
    result
}

add then looks just like we'd like it to:

fn add(a: Coord, b: Coord) -> Coord {
    new_coord_from(a.iter().zip(&b).map(|(a, b)| a + b))
}

Upvotes: 7

Related Questions