如何无需循环通过索引数组从三维numpy数组取值填充二维数组

你可以直接使用numpy的向量化索引实现，无需任何Python层循环，性能远高于嵌套循环实现：

方法1：手动构造高级索引 #

通过np.indices生成前两维的索引数组，和b组合直接索引a：

import numpy as np

# 你的原始数组生成代码
m, n, k  = 3, 2, 4
np.random.seed(1234)
a = np.random.randint(0,9,(m, n, k))
b = np.random.randint(0,k,(m,n))

# 核心实现
i, j = np.indices(b.shape)
c = a[i, j, b]

print(c)
# 输出：
# [[4 8]
#  [0 5]
#  [7 2]]

方法2：使用take_along_axisAPI（更简洁） #

numpy内置的take_along_axis专门适配这种场景，不需要手动构造前序维度的索引：

# 核心实现，axis=2指定沿第三维取值
c = np.take_along_axis(a, b[..., np.newaxis], axis=2).squeeze(axis=2)

两种实现的时间复杂度均为O(mn)，完全由numpy底层C代码执行，哪怕是千万级元素的超大数组也可以高效运行。

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