如何从二维numpy数组中选择样本？

Question

我需要根据列表在2D NumPy矩阵中选择一些单元格。但是，以下代码

a = np.array([[1,2,3],[4,5,6],[7,8,9]])
b = [(1,1), (1,2), (2,1)]
a[b]

引发以下错误：

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-195-a8bd6862e58c> in <module>()
      1 a = np.array([[1,2,3],[4,5,6],[7,8,9]])
      2 b = [(1,1), (1,2), (2,1)]
----> 3 a[b]

IndexError: too many indices for array

但是，正确的结果必须是(5, 6, 8)。我可以使用下面的代码完成它，但是，我想知道是否有更有效的方法（以numpy矢量化方式）？

有效的代码：

map(lambda (x): a[x], b)

提前致谢。

Answer 1

使用b

转置 zip，然后将其用作索引：

a[tuple(zip(*b))]
# array([5, 6, 8])

相当于：

row, col = zip(*b)
a[row, col]