python - 在Numpy中逐行索引

我有两个矩阵，A和B：

A = array([[2., 13., 25., 1.], [ 18., 5., 1., 25.]])
B = array([[2, 1], [0, 3]])

我想用每一行的A索引每一行，生成切片：

array([[25., 13.], [18., 25.]])

也就是说，我基本上想要：

array([A[i,b] for i,b in enumerate(B)])

有没有办法直接用花式索引？我能做的最好的就是这个“平头黑客”：

A.flat[B + arange(0,A.size,A.shape[1])[:,None]]

最佳答案

@奥菲翁的回答很好，值得赞扬，但我想补充一些解释，并提供一个更直观的结构。
与其先旋转B然后再旋转结果，最好只旋转arange即可。我认为这是最直观的解决方案，即使需要更多的字符：

A[((0,),(1,)), B]

或同等

A[np.arange(2)[:, None], B]

这是因为这里真正发生的事情是，你制作了一个i数组和一个j数组，每个数组的形状都与你想要的结果相同。

i = np.array([[0, 0],
              [1, 1]])
j = B

但是你可以用

i = np.array([[0],
              [1]])

因为它将广播以匹配B（这是np.arange(2)[:,None]给出的）。
最后，为了使其更通用（不知道2是arange大小），您还可以使用

i = np.indices(B.shape)[0]

不管你如何构建i和B，你只需要把它称为

>>> A[i, j]
array([[ 25.,  13.],
       [ 18.,  25.]])