python - 如何从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.]])