通用数组填充算法实现：支持动态指定维度

我有一段处理3维数组的代码，但希望能动态指定维度。Matlab的nchoosek函数生成的是无重复元素的数组，不符合我的需求。

我的3维实现代码如下：

energy_band=70;
count=1;
for i=-energy_band:energy_band
    for j=-energy_band:energy_band
        for z=-energy_band:energy_band
          out(1,count)=i;
          out(2,count)=j;
          out(3,count)=z;
          count=count+1;
       end
   end
end

4维的实现类似，仅需多套一层循环：

for i=-energy_band:energy_band % 能量偏移范围从-1.75 eV到1.75 eV的数组生成循环
            for j=-energy_band:energy_band
                for z=-energy_band:energy_band
                    for k=-energy_band:energy_band
                        out(1,count)=i;
                        out(2,count)=j;
                        out(3,count)=z;
                        out(4,count)=k;
                        count=count+1;
                    end
                end
            end
        end

目前我通过switch枚举维度的方式实现，但这种方法缺乏通用性，新增维度就得手动添加case：

function out=energy_array(matrix_version,energy_band)
count=1;
switch matrix_version
    case 1
        for i=-energy_band:energy_band % 能量偏移范围从-1.75 eV到1.75 eV的数组生成循环
            out(1,count)=i;
            count=count+1;
        end
    case 2
        for i=-energy_band:energy_band % 能量偏移范围从-1.75 eV到1.75 eV的数组生成循环
            for j=-energy_band:energy_band
                out(1,count)=i;
                out(2,count)=j;
                count=count+1;
            end
        end
    case 3
        for i=-energy_band:energy_band % 能量偏移范围从-1.75 eV到1.75 eV的数组生成循环
            for j=-energy_band:energy_band
                for z=-energy_band:energy_band
                    out(1,count)=i;
                    out(2,count)=j;
                    out(3,count)=z;
                    count=count+1;
                end
            end
        end
    case 4
        for i=-energy_band:energy_band % 能量偏移范围从-1.75 eV到1.75 eV的数组生成循环
            for j=-energy_band:energy_band
                for z=-energy_band:energy_band
                    for k=-energy_band:energy_band
                        out(1,count)=i;
                        out(2,count)=j;
                        out(3,count)=z;
                        out(4,count)=k;
                        count=count+1;
                    end
                end
            end
        end
    case 5
        for i=-energy_band:energy_band % 能量偏移范围从-1.75 eV到1.75 eV的数组生成循环
            for j=-energy_band:energy_band
                for z=-energy_band:energy_band
                    for k=-energy_band:energy_band
                        for l=-energy_band:energy_band
                            out(1,count)=i;
                            out(2,count)=j;
                            out(3,count)=z;
                            out(4,count)=k;
                            out(5,count)=l;
                            count=count+1;
                        end
                    end
                end
            end
        end
end

请问如何更优雅地实现该功能，无需枚举所有维度情况？

方法1：利用ndgrid+reshape实现（最简洁高效） #

Matlab的ndgrid可生成任意维度的网格坐标，结合reshape和cellfun就能快速得到目标格式的结果，完全无需手动嵌套循环：

function out = energy_array(dim_num, energy_band)
    % 生成单个维度的取值范围
    range = -energy_band : energy_band;
    % 生成ndgrid的输入参数，重复range指定的维度次数
    grid_args = repmat({range}, 1, dim_num);
    % 调用ndgrid生成多维度网格
    [grid_vars{1:dim_num}] = ndgrid(grid_args{:});
    % 将每个网格变量转置后拉成列向量，再拼接为目标格式
    out = cellfun(@(x) reshape(x.', [], 1), grid_vars, 'UniformOutput', false);
    out = [out{:}].';
end

说明： #

• ndgrid(grid_args{:})会生成dim_num个网格矩阵，每个矩阵对应一个维度的坐标值
• 最终输出的矩阵格式和你原有代码完全一致：行对应维度，列对应每个坐标组合

方法2：递归实现动态嵌套循环 #

如果更倾向于循环逻辑，可通过递归动态生成嵌套循环，避免枚举维度：

function out = energy_array(dim_num, energy_band)
    range = -energy_band : energy_band;
    total_samples = length(range)^dim_num;
    out = zeros(dim_num, total_samples);
    % 调用递归函数填充结果
    recursive_fill(out, range, dim_num, 1, []);
end

function recursive_fill(out, range, remaining_dims, current_col, current_vals)
    if remaining_dims == 0
        % 所有维度遍历完成，填充当前列
        out(:, current_col) = current_vals;
        return;
    end
    % 遍历当前维度的所有取值
    for val = range
        new_vals = [current_vals; val];
        recursive_fill(out, range, remaining_dims - 1, current_col, new_vals);
        current_col = current_col + 1;
    end
end

说明： #

• 递归函数通过逐层减少剩余维度数，模拟嵌套循环的逻辑，当剩余维度为0时，将当前坐标组合填充到结果矩阵的对应列

方法3：手动生成笛卡尔积 #

也可通过迭代生成笛卡尔积的方式实现，逻辑直观：

function out = energy_array(dim_num, energy_band)
    range = -energy_band : energy_band;
    % 初始化1维笛卡尔积
    cart_prod = range.';
    % 迭代生成更高维度的笛卡尔积
    for d = 2:dim_num
        [a, b] = meshgrid(cart_prod, range);
        cart_prod = [reshape(a, [], 1), reshape(b, [], 1)];
    end
    % 转置得到目标格式
    out = cart_prod.';
end

说明： #

• 从1维开始，每次迭代将当前笛卡尔积与新维度的取值做meshgrid，再拼接成新的笛卡尔积矩阵，最后转置得到目标格式

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