如何实现N维NumPy数组切片的遍历求和泛化？

问题背景 #

你已经将多维度的pandas DataFrame转换为N维NumPy数组，希望遍历每个维度的所有切片并求和，但原有代码无法适配任意维度（比如4维数组需要手动写arr[ii,:,:,:]），需要一种通用的实现方式。

核心思路 #

NumPy支持用...（省略号）表示“剩余所有维度的全部元素”，或者用slice(None)动态构建索引元组，这两种方式都能实现任意维度的切片泛化，无需硬编码维度数量。

方案1：使用省略号...简化切片 #

直接用arr[ii, ...]代替固定数量的冒号，...会自动匹配剩余的所有维度，代码无需修改维度相关的硬编码：

import numpy as np
import pandas as pd

# 示例DataFrame
df = pd.DataFrame({'dim1': ['a', 'a', 'b', 'b'], 'dim2': ['x', 'y', 'x', 'y'], 'val': [2, 4, 6, 8]})

# 转换函数保持不变
def df_to_numpy(df: pd.DataFrame) -> np.array:
    try:
        shape = [len(level) for level in df.index.levels]
    except AttributeError:
        shape = [len(df.index)]
    ncol = df.shape[-1]
    if ncol > 1:
        shape.append(ncol)
    return df.to_numpy().reshape(shape)

# 转换为N维数组（可扩展到任意维）
arr = df_to_numpy(df.set_index(['dim1', 'dim2']).unstack())

# 泛化遍历每个维度的切片求和
for dim_idx in range(arr.ndim):
    # 将当前维度移到第一个位置，统一处理逻辑
    arr_shifted = np.moveaxis(arr, dim_idx, 0)
    print(f"=== 处理第 {dim_idx} 维度 ===")
    for ii in range(arr_shifted.shape[0]):
        print(f"唯一分组编号: {ii}")
        print(f"求和结果: {arr_shifted[ii, ...].sum()}")

方案2：动态构建索引元组 #

如果需要更灵活的索引控制，可以用slice(None)构建索引元组，效果和...一致：

for dim_idx in range(arr.ndim):
    arr_shifted = np.moveaxis(arr, dim_idx, 0)
    print(f"=== 处理第 {dim_idx} 维度 ===")
    # 构建索引模板：(ii,) + 剩余维度的slice(None)
    for ii in range(arr_shifted.shape[0]):
        idx = (ii,) + (slice(None),) * (arr_shifted.ndim - 1)
        print(f"唯一分组编号: {ii}")
        print(f"求和结果: {arr_shifted[idx].sum()}")

方案3：利用NumPy的sum指定axis（更高效） #

如果不需要逐个打印切片，只是需要计算每个维度所有切片的求和结果，可以直接用np.sum并指定求和的轴，避免手动循环：

for dim_idx in range(arr.ndim):
    # 计算除当前维度外的所有轴的和
    sum_axes = tuple(i for i in range(arr.ndim) if i != dim_idx)
    dim_sums = arr.sum(axis=sum_axes)
    print(f"第 {dim_idx} 维度各切片求和结果: {dim_sums}")

原有代码无法泛化的原因 #

原有代码用swapaxes交换轴后，硬编码了arr[ii,:]，当维度超过2时，:只能匹配一个维度，剩余维度没有被覆盖，导致切片范围错误。而用...或动态索引元组可以自动适配任意数量的剩余维度。

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