使用NumPy生成C顺序N维索引数组的高效可读方法

我需要生成C顺序的N维索引数组，尝试了三种NumPy实现方式：

• 大数组速度最快但可读性差：

  np.stack(np.meshgrid(*[np.arange(i, dtype=dtype) for i in sizes], indexing="ij"), axis=-1).reshape(-1, len(sizes))
  

• 可读性好且性能优异：

  np.ascontiguousarray(np.indices(sizes, dtype=dtype).reshape(len(sizes), -1).T)
  

  这里不确定是否必须用ascontiguousarray强制拷贝，有没有更好的方式在不拷贝的前提下保证结果是C连续的。
• 可读性最佳但速度最慢：

  np.vstack([*np.ndindex(sizes)], dtype=dtype)
  

  大数组场景下，迭代器转换的速度很慢。

请问NumPy有没有内置更直观易读、性能媲美np.meshgrid或np.indices的实现？如果没有，能不能优化meshgrid或indices的方案，在避免不必要内存拷贝（比如ascontiguousarray）的同时确保数组是C连续的？

示例代码：

sizes = (3, 1, 2)
idx = np.ascontiguousarray(np.indices(sizes).reshape(len(sizes), -1).T)
print(idx)
print(f"C_CONTIGUOUS: {idx.flags['C_CONTIGUOUS']}")
# 输出：
# [[0 0 0]
#  [0 0 1]
#  [1 0 0]
#  [1 0 1]
#  [2 0 0]
#  [2 0 1]]
# C_CONTIGUOUS: True

解答 #

关于NumPy内置方案 #

目前NumPy没有专门的内置函数直接生成直观易读且性能拉满的C顺序N维索引数组，但可以通过优化现有indices或meshgrid的用法来达到目标。

优化np.indices方案（避免不必要拷贝） #

np.indices默认返回形状为(N, *sizes)的数组，每个子数组对应一个维度的索引，且默认是Fortran连续布局。直接转置再重塑会导致数组非连续，因此需要ascontiguousarray拷贝，但可以通过调整维度顺序的方式避免强制拷贝：

def c_order_indices(sizes, dtype=None):
    idx = np.indices(sizes, dtype=dtype)
    # 将维度从 (N, d1, d2, ..., dn) 转成 (d1, d2, ..., dn, N)，再展平为 (总元素数, N)
    return idx.transpose(*range(1, len(sizes)+1), 0).reshape(-1, len(sizes))

测试验证：

sizes = (3,1,2)
idx = c_order_indices(sizes)
print(idx.flags['C_CONTIGUOUS'])  # 输出 True
print(idx)
# [[0 0 0]
#  [0 0 1]
#  [1 0 0]
#  [1 0 1]
#  [2 0 0]
#  [2 0 1]]

该方案的核心逻辑是：将索引维度移到最后后，数组的内存布局天然符合C连续规则，后续的reshape操作不会改变内存连续性（只要形状兼容），因此无需额外拷贝。

性能与可读性对比 #

• 性能：和np.indices原生实现几乎一致，完全媲美meshgrid方案
• 可读性：逻辑清晰，直接体现“生成索引→调整维度顺序→展平为每行一个索引”的流程，比meshgrid写法更简洁

关于np.meshgrid的优化 #

如果偏好meshgrid的写法，可以简化为：

np.stack(np.meshgrid(*map(np.arange, sizes), indexing="ij"), axis=-1).reshape(-1, len(sizes))

该写法性能和优化后的indices方案接近，但可读性稍弱，需要手动生成维度序列并执行stack操作。

总结 #

最优方案是采用调整维度顺序的np.indices写法，既保证了C连续布局、无不必要内存拷贝，又兼顾了可读性和性能，完全满足需求。

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