Rust ndarray布尔掩码索引效率优化：求高效实现方案

问题背景 #

需要在Rust ndarray中实现类似Numpy的分组均值计算，原实现通过布尔掩码逐组筛选数据，效率极低（每次迭代耗时接近Numpy整个脚本的执行时间）。

Numpy参考实现（高效） #

import numpy as np

shape = (100, 100, 100)

grouping_array = np.random.randint(0, 100, size=shape)
data_array = np.random.rand(*shape)

for i in range(1, 100):
    ith_mean = data_array[grouping_array == i].mean()
    print(ith_mean)

原Rust实现（低效） #

fn group_means(
    data: &Array<f32, IxDyn>,
    grouping_var: &Array<f32, IxDyn>,
    n_groups: i32,
) {
    
    for group in 1..n_groups {
        
        let index_array = grouping_var.mapv(|x| x == group as f32); // 原代码中roi应为group
        let roi_data = Array::from_iter(
            data // 原代码中image_data应为data
            .iter()
            .zip(index_array.iter())
            .map(|(x, y)| if *y { *x } else { 0. })
        );
        
        let mean_roi = roi_data.mean().unwrap();
        println!("group {}; mean {}", group, mean_roi);
    }
    
}

优化方案：单次遍历统计所有组 #

原实现的核心问题是每次循环都要遍历整个数组并创建新数组，时间复杂度为O(k*n)（k为组数，n为元素总数）。优化思路改为单次遍历完成所有组的总和与计数统计，时间复杂度降为O(n)，同时避免冗余内存操作。

优化后的Rust代码 #

use ndarray::Array;
use ndarray::IxDyn;

fn group_means(
    data: &Array<f32, IxDyn>,
    grouping_var: &Array<i32, IxDyn>, // 改用整数类型存储分组，避免浮点数精度问题
    n_groups: i32,
) {
    // 初始化每个组的总和与计数，索引对应组号
    let mut sums = vec![0.0f32; n_groups as usize];
    let mut counts = vec![0usize; n_groups as usize];

    // 一次遍历完成所有组的统计
    data.iter().zip(grouping_var.iter()).for_each(|(&val, &group)| {
        let group_idx = group as usize;
        // 确保组号在有效范围内
        if group >= 0 && group_idx < n_groups as usize {
            sums[group_idx] += val;
            counts[group_idx] += 1;
        }
    });

    // 计算并打印每个组的均值（从1开始）
    for group in 1..n_groups {
        let idx = group as usize;
        if counts[idx] > 0 {
            let mean = sums[idx] / counts[idx] as f32;
            println!("group {}; mean {}", group, mean);
        } else {
            println!("group {}; no elements", group);
        }
    }
}

优化点说明 #

• 单次遍历统计：仅遍历数据和分组数组一次，将每个元素的数值累加到对应组的总和，同时计数，避免了原代码中k次重复遍历的开销
• 消除冗余内存操作：原代码每次循环都会创建新的掩码数组和筛选后的数据数组，涉及大量内存分配与元素拷贝；优化后仅用两个Vec存储中间结果，内存开销极小
• 改用整数分组：避免浮点数相等判断的精度问题，同时整数类型的索引访问比浮点数比较更高效
• 边界检查：增加组号的有效性判断，避免数组越界访问

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