如何使用每行的索引矩阵来索引矩阵的行元素？答案

【问题标题】：How to index row elements of a Matrix with a Matrix of indices for each row?如何使用每行的索引矩阵来索引矩阵的行元素？
【发布时间】：2021-06-17 18:09:53
【问题描述】：

我有一个索引矩阵，例如

I = np.array([[1, 0, 2], [2, 1, 0]])

第 i 行的索引从第 i 行的另一个矩阵 M 中选择一个元素。

所以有我，例如

M = np.array([[6, 7, 8], [9, 10, 11])

M[I] 应该选择：

[[7, 6, 8], [11, 10, 9]]

我可以：

I1 = np.repeat(np.arange(0, I.shape[0]), I.shape[1])
I2 = np.ravel(I)
Result = M[I1, I2].reshape(I.shape)

但这看起来很复杂，我正在寻找更优雅的解决方案。最好不要压扁和整形。

在示例中我使用了 numpy，但实际上我使用的是 jax。所以如果jax中有更高效的解决方案，欢迎分享。

【问题讨论】：

标签： python numpy matrix-indexing

【解决方案1】：

In [108]: I = np.array([[1, 0, 2], [2, 1, 0]])
     ...: M = np.array([[6, 7, 8], [9, 10, 11]])
     ...: 
     ...: I,M

我必须在 M 中添加一个 ']'。

Out[108]: 
(array([[1, 0, 2],
        [2, 1, 0]]),
 array([[ 6,  7,  8],
        [ 9, 10, 11]]))

broadcasting 的高级索引：

In [110]: M[np.arange(2)[:,None],I]
Out[110]: 
array([[ 7,  6,  8],
       [11, 10,  9]])

第一个索引的形状为 (2,1)，与 I 的 (2,3) 形状配对以选择 (2,3) 值块。

【讨论】：

【解决方案2】：

这一行代码怎么样？这个想法是枚举矩阵的行和行索引，因此您可以访问索引矩阵中的相应行。

import numpy as np

I = np.array([[1, 0, 2], [2, 1, 0]])
M = np.array([[6, 7, 8], [9, 10, 11]])

Result = np.array([row[I[i]] for i, row in enumerate(M)])
print(Result)

输出：

[[ 7  6  8]
 [11 10  9]]

【讨论】：

您可以在没有列表理解的情况下执行此操作。

【解决方案3】：

np.take_along_axis 也可以在这里使用I 超过axis=1 的索引来获取M 的值：

>>> np.take_along_axis(M, I, axis=1)

array([[ 7,  6,  8],
       [11, 10,  9]])

【讨论】：