python - 由其他向量确定的多个轴上的 Numpy argmax

Question

假设我有一个名为 A 的 3d 矩阵，其形状为 (100,4,100)。还有两个向量，每个向量有 200 个值，创建如下：

b = np.random.randint(0, 3, 200)
c = np.random.randint(0, 99, 200)

如何有效地找到每个 A[b,c,:] 的 argmax？

我可以做一个 for 循环，看起来像这样：

for i in range(200):
    np.argmax(A[b[i], c[i], :])

通过避免 for 循环，我是否有更有效的选择？

Answer 1

上面 Divakar 建议的解决方案是使用ndarray.argmax，所以它看起来像A[b,c].argmax(1)