Rust中高效处理ndarray：单次遍历生成双数组，避免重复计算与分配

需求与现有实现 #

希望遍历ndarray的元素并进行转换，同时单次计算得到两个新结果，后续基于这两个结果计算点积。现有Rust伪代码如下：

use ndarray::*; // definition of Array1

// let's assume examplary input/output
let x = Array1::random(100, Uniform::<f64>::new(0., 1.));
let y = Array1::random(100, Uniform::<f64>::new(0., 1.));

// transformation of the x array into two arrays (a,b) of the same size
let (a, b): (Array1<f64>, Array1<f64>) = x
    .iter()
    .map(|v| {
        // ...  computing
        let result1 = v; // just dummy example
        let result2 = v*2.0; // just dummy example
        return (result1, result2);
    })
    .collect(); // or unzip()

// after this I’d like to compute something on the result
let res1: f64 = a.dot(&y);
let res2: f64 = b.dot(&y);

现有痛点 #

• 转换为Vec再转Array会产生额外的内存分配与转换开销，后续点积计算希望避免此问题；
• 拆分闭包用mapv_inplace等方法会重复执行昂贵的计算逻辑，不可接受。

解决方案 #

方案一：直接计算点积，跳过中间数组（最优选择） #

如果后续只需要点积结果，完全可以跳过中间数组a和b的生成，在单次遍历中同时计算两个点积的累加值，既节省内存又避免重复计算：

use ndarray::*;
use rand::distributions::Uniform;
use rand::Rng;

fn main() {
    let mut rng = rand::thread_rng();
    let x = Array1::random_using(100, Uniform::new(0., 1.), &mut rng);
    let y = Array1::random_using(100, Uniform::new(0., 1.), &mut rng);

    // 初始化点积结果的累加器
    let (mut res1, mut res2) = (0.0f64, 0.0f64);

    // 同时遍历x和y的对应元素
    for (x_val, y_val) in x.iter().zip(y.iter()) {
        // 执行一次昂贵的计算逻辑，得到两个结果
        let result1 = *x_val; // 替换为实际计算逻辑
        let result2 = *x_val * 2.0; // 替换为实际计算逻辑

        // 直接累加点积项
        res1 += result1 * y_val;
        res2 += result2 * y_val;
    }

    println!("res1: {}, res2: {}", res1, res2);
}

优势：

• 完全消除中间数组的内存分配，降低内存占用；
• 昂贵计算逻辑仅执行一次，避免重复开销；
• 时间复杂度保持O(n)，但减少了中间步骤的额外耗时。

方案二：预分配数组并通过Zip填充（需保留中间数组时使用） #

如果后续还需要使用a和b数组，可以通过预分配固定大小的数组，再用ndarray的Zip工具实现一次遍历同时填充两个数组：

use ndarray::*;
use rand::distributions::Uniform;
use rand::Rng;

fn main() {
    let mut rng = rand::thread_rng();
    let x = Array1::random_using(100, Uniform::new(0., 1.), &mut rng);
    let y = Array1::random_using(100, Uniform::new(0., 1.), &mut rng);

    // 预分配与x同大小的空数组，避免动态扩容开销
    let mut a = Array1::zeros(x.len());
    let mut b = Array1::zeros(x.len());

    // 使用Zip同时遍历x并写入a、b数组
    Zip::from(&mut a)
        .and(&mut b)
        .and(&x)
        .apply(|a_val, b_val, x_val| {
            // 仅执行一次昂贵计算
            let result1 = *x_val; // 替换为实际计算逻辑
            let result2 = *x_val * 2.0; // 替换为实际计算逻辑
            *a_val = result1;
            *b_val = result2;
        });

    // 后续正常计算点积
    let res1 = a.dot(&y);
    let res2 = b.dot(&y);

    println!("res1: {}, res2: {}", res1, res2);
}

优势：

• 预分配数组避免了collect/unzip带来的Vec转Array的额外转换；
• 通过Zip实现单次遍历填充两个数组，保证计算逻辑只执行一次；
• 后续操作可直接使用预分配的数组，和原逻辑兼容。

内容的提问来源于stack exchange，提问作者mathmax