如何逐元素地将两个 Rust 数组相加?

9

这是一个非常初级的问题,但我在搜索了半个小时后没有找到有用的信息。我有Rust 1.7.0版本和以下代码:

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 {
    //???
}

显而易见我首先尝试的是a.zip(b).map(|(u,v)| u+v),但这并不起作用(无法压缩阵列),也不会用a.iter().zip(b.iter()).map(|(u,v)| u+v),因为它无法将迭代器转换回数组。我可以理解为什么这在一般情况下不起作用,但在这种情况下,我们知道两个变量具有相同的长度。
目前我正在进行以下操作:
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
}

但与之相比,它看起来很丑。我错过了什么吗?

5
请注意,在您的代码中,0..2 应该更改为 0..3 - wimh
2
“丑陋”是主观的,但我认为显式的 for i in 0..3 比提出的其他选择更易读。 - Paolo Falabella
1个回答

6

一个简单的方法是使用enumerate迭代器方法生成索引,并通过将值分配到获得的索引中来填充z

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]);
}

在Rust中,我们可以比此更好地处理它,通过注意到iter()生成的是数组中的引用,并且有iter_mut()可用于生成可变引用。这会导致代码非常类似于您尝试编写的内容:
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
}

如果这种写入z的模式会反复出现,并带有不同的操作,您可以将新Coord的创建和数据填充抽象成通用函数:

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 现在看起来就像我们想要的一样:

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

小的后续问题:如果我想声明一个函数 fn cmap(a: Coord, b: Coord, f: F) -> Coord { new_coord_from(a.iter().zip(&b).map(f)); } 来抽象出迭代/压缩部分,以便我可以声明 fn add(a: Coord, b: Coord) -> Coord { cmap(a,b, |(a,b)| a + b) } - 我需要为 F 放置什么类型? - user1569030
1
目前的写法是 cmap<F: Fn((&i64, &i64)) -> i64>(a: Coord, b: Coord, f: F) -> Coord。然而,我更喜欢去掉引用,并让 F 接收两个数字而不是元组:fn cmap<F: Fn(i64, i64) -> i64>(a: Coord, b: Coord, f: F) -> Coord { new_coord_from(a.iter().zip(&b).map(|(x, y)| f(*x, *y))) }。然后,add 的定义就简化为最简单的形式,即 cmap(a, b, |x, y| x + y) 或者甚至是 cmap(a, b, i64::add)(在使用 use std::ops::AddAdd 特性引入作用域之后)。 - user4815162342

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接