如何在python中从nd数组中提取特定行

Question

我有一个数组b，它的形状是E x B x 3。我还有另一个数组a，它指定了取b的3个元素

以下代码起作用（在本例中E=2, B=4）：

import numpy as np

a = [1, 1, 0, 0]
b = np.array([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 3.0]],
            [[2.0, 0.0, 0.0], [0.0, 2.0, 0.0],  [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]]])
# n_pred = np.transpose(n_pred, axes=[1, 0, 2])
c = []
for i, idx in enumerate(a):
    c.append(b[idx, i])
c = np.array(c)
print(c)

我的问题是，有没有更有效的方法(也许是使用了一些内置的numpy函数

Answer 1

可以按前两个维度进行索引：

c = b[a, range(len(a))]

print(c)

array([[ 2.,  0.,  0.],
       [ 0.,  2.,  0.],
       [ 0.,  0.,  1.],
       [ 0.,  0.,  3.]])