Numpy 中的逐行索引

Posted 2023-02-23

【中文标题】Numpy 中的逐行索引

【英文标题】：Row-wise indexing in Numpy

【发布时间】：2013-11-15 13:12:26

【问题描述】：

我有两个矩阵，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)])

有没有办法直接对它进行幻想索引？我能做的最好的就是这个“flat-hack”：

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

【问题讨论】：

另见***.com/questions/10921893/…

【参考方案1】：

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

与其先旋转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），您还可以从B 生成i

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

无论您如何构建 i 和 j，都可以这样称呼它

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

【讨论】：

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

使用np.r_ 也有一些很棒的解决方案，但这真的让我很困惑，并且是为习惯MATLAB 表示法的人设计的，我不是。

@askewchan，我在这里解决了这个问题的 N 维版本：***.com/questions/10921893/…

【参考方案2】：

不漂亮但是：

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

【讨论】：

啊哈，我没想过要转置B。这样更干净。