如何使用1d位置数组切割numpy数组

Question

考虑数组a和b

a = np.arange(25).reshape(5, -1)
b = np.array([4, 2, 3, 0, 1])

如何使用a对b进行切片以获取这些元素？

我可以使用列表理解

np.array([a[i, b[i]] for i in range(5)])

array([ 4,  7, 13, 15, 21])

但我想要一个纯粹的numpy解决方案

Answer 1

试试这个：

>>> a[np.arange(5), b]
array([ 4,  7, 13, 15, 21])

当索引是数组时，它们在broadcasting rules之后按元素解释。