NumPy中的按行索引

Question

我有两个矩阵，A和B：

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

我希望用B的每一行去索引A的每一行，生成切片：

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

也就是说，我基本上想要像这样的东西：

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

有没有直接进行“fancy-index”（精美索引）的方法？我所能做到的最好的方法是这个“flat-hack”：

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

- perimosocordiae

---

请参见https://dev59.com/fGgu5IYBdhLWcg3w_cBV#10922358。 - user545424

Answer 1

@Ophion的回答很好，值得赞扬，但我想添加一些解释，并提供更直观的构造。

不要旋转B然后将结果旋转回来，最好只旋转arange。我认为这提供了最直观的解决方案，即使需要更多字符：

而不是旋转 B 并将结果旋转回来，最好只旋转 arange。我认为这是最直观的解决方案，即使需要更多的字符：

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

或等价于

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

这样做的原因是因为实际上你正在创建一个数组和一个数组，它们的形状与你想要的结果相同。

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

但是你可以只使用

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

因为它将广播到与 B 相匹配（这就是 np.arange(2)[:,None] 的作用）。

最后，为了使其更加通用（不知道 arange 大小的 2），您还可以使用 B 生成 i：

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

无论你如何构建i和j，你只需要调用它。

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

- askewchan

---

这是更加深思熟虑的解决方案。我特别喜欢它，因为A.flat[np.ravel_multi_index((np.arange(2)[:,None],B),A.shape)]这样展示更有意义。 - Daniel

还有一些很好的解决方案使用 np.r_，但那真的让我感到困惑，并且它是为了习惯于 MATLAB 表示法的人设计的，而我并不是。 - askewchan

@askewchan，我在这里解决了这个问题的N维版本：https://dev59.com/fGgu5IYBdhLWcg3w_cBV#10922358 - user545424

Answer 2

虽然不是很美观，但是：

A[np.arange(2),B.T].T
array([[ 25.,  13.],
       [ 18.,  25.]])

- Daniel

---

啊哈，我没想到要转置B。这样更简洁。 - perimosocordiae