如何无嵌套循环高效用2D矩阵C切片3D矩阵A？

这问题我太熟了，numpy里处理这种“按索引取对应位置元素”的场景，有两个超好用的无循环方案，比你说的布尔掩码高效多了，而且代码简洁：

方法1：直接用广播式索引（最简洁） #

numpy支持多维数组的索引数组广播，我们只需要生成和C对应位置的y、x坐标，再和C一起作为索引传入A就行：

import numpy as np

# 先生成y和x的网格索引数组（和C形状完全一致）
y_idx, x_idx = np.indices(C.shape)
# 直接通过三维索引提取值，numpy会自动对齐维度
D = A[C, y_idx, x_idx]

原理很简单：np.indices(C.shape)会输出两个(h,w)的数组，y_idx[y,x] = y，x_idx[y,x] = x，把C作为z轴的索引，三个数组组合起来就精准定位到了A中每个(C[y,x], y, x)的元素，完全替代嵌套循环，速度是循环的几十上百倍。

方法2：用np.take_along_axis（更直观的轴对齐） #

如果你觉得索引广播有点绕，np.take_along_axis这个函数就是专门干“沿指定轴按索引取元素”的活的，用法也很清晰：

# 把C扩展成和A的轴0匹配的形状（增加一个前置维度，变成(1, h, w)）
C_expanded = C[np.newaxis, :, :]
# 沿轴0（z轴）提取对应索引的元素，再压缩掉多余的维度
D = np.take_along_axis(A, C_expanded, axis=0).squeeze(axis=0)

这个函数的好处是逻辑更直白：明确告诉numpy“沿着z轴，每个位置取C指定的索引值”，适合对轴操作更敏感的场景。

为啥不推荐布尔掩码？ #

你之前想的布尔掩码方法确实可行，但会生成一个和A一样大的(n,h,w)布尔数组，内存占用比上面两种方法高很多，而且计算效率也不如直接索引——毕竟索引是直接定位元素，掩码还要做乘法和求和，完全没必要绕这个弯。

另外你提到的np.choose确实有元素数量限制（老版本numpy限制32，新版本放宽了但还是有局限），所以上面两个方法才是生产环境里的首选。

用你给的示例数据测试的话，这两种方法得到的D和循环版完全一致，放心用就行~

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