GPU数组兼容的Matlab多占位find函数向量化改造需求

问题背景 #

我们需要将N个粒子在M个格子中的状态向量（如[2 1 3 0]表示盒子1有2个粒子、盒子2有1个、盒子3有3个、盒子4空）转换为粒子位置向量（对应输出[1 1 2 3 3 3]）。原实现使用了循环和cell数组，当输入为gpuArray时，cell2mat不兼容，因此需要移除循环与cell数组，实现纯向量化代码。

向量化解决方案 #

可以利用Matlab内置的repelem函数（R2015a及以上版本支持，且完全兼容GPU数组）直接实现需求，同时简化原函数的分支逻辑：

function position = findGeneral(state, N)
    % 获取非零粒子数的盒子索引
    box_indices = find(state);
    % 获取对应盒子的粒子数量
    particle_counts = state(box_indices);
    % 按粒子数重复盒子索引，生成位置向量
    position = repelem(box_indices, particle_counts);
    
    % 确保输出长度严格为N（处理状态向量总和与N不一致的边界情况）
    pos_length = numel(position);
    if pos_length > N
        position = position(1:N);
    elseif pos_length < N
        position = [position, zeros(1, N - pos_length, class(state))];
    end
end

方案说明 #

• GPU兼容性：repelem原生支持gpuArray输入，无需任何cell操作，彻底解决原代码的兼容性问题。
• 逻辑简化：去掉了原函数中针对"是否有粒子数≥2的盒子"的分支判断，无论状态向量是全1还是包含多粒子盒子，都能统一处理。
• 边界处理：添加了长度校验，确保输出向量长度严格等于N，避免因状态向量总和与N不匹配导致的错误。

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