NumPy中如何不使用for循环，以数组为索引定义新矩阵？

问题描述 #

• 现有维度为(xsiz, zsiz)的x-z平面数组xz，以及维度为ysiz的y方向索引数组y，初始化代码如下：

xz = np.zeros((xsiz, zsiz))
y = (np.arange(ysiz)*(zsiz/ysiz)).astype(int)

• xz可看作由zsiz个长度为xsiz的列向量组成，每个列向量对应[0, zsiz-1]范围内的一个z值标记，此前提取单z对应列的写法为np.transpose(xz)[z]。
• 需求：以y数组作为z值的索引列表，提取对应z值标记的列向量，最终拼接为维度(xsiz, ysiz)的x-y平面矩阵。
• 约束：不使用for循环，全程通过数组原生操作实现，以适配CuPy的GPU并行加速要求。

参考效果的for循环实现 #

以下为可正确得到目标结果的循环代码，用于对齐需求：

import numpy as np

xsiz = 5 # 示例用测试尺寸
ysiz = 6
zsiz = 4

xz = np.arange(xsiz*zsiz).reshape(xsiz, zsiz)
y = (np.arange(ysiz)*(zsiz/ysiz)).astype(int)

xzT = np.transpose(xz)

final_xyT = np.zeros((0, xsiz))

for i in range(ysiz):
    index = y[i]
    xvec = xzT[index]
    final_xyT = np.vstack((final_xyT, xvec))

# 若y中存在超出z轴范围的索引可能触发报错
# CuPy索引默认支持环绕取值，该问题影响较小

final_xy = np.transpose(final_xyT)
print(xz)
print(final_xy)

无循环向量化实现 #

你猜测的「转置后索引再转置」思路完全可行，利用NumPy/CuPy原生的高级索引功能可以一次性完成取数，不需要任何循环，两种等价写法如下：

# 写法1：沿你之前的思路，转置后按行索引再转置回来
final_xy = xz.T[y].T

# 写法2：直接对z轴维度做索引，无需转置，逻辑更直观
final_xy = xz[:, y]

两种写法输出结果和for循环完全一致：写法2中[:, y]表示保留x轴（第一维）的所有元素，z轴（第二维）按照y数组给出的索引顺序选取元素，输出维度直接为(xsiz, ysiz)。
该实现为纯向量化数组操作，无Python层面的循环开销，在CuPy中可自动调度GPU完成并行计算，性能远高于逐次vstack的循环实现。

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