从索引矩阵向量化数组访问答案

【问题标题】：Vectorizing array access from indices matrix从索引矩阵向量化数组访问
【发布时间】：2019-06-14 22:34:15
【问题描述】：

考虑以下几点：

In [51]: arr = np.arange(6, 10)

In [52]: idx = np.random.randint(4, size=(3, 4))

In [53]: idx
Out[53]:
array([[0, 3, 3, 1],
    [1, 3, 3, 2],
    [1, 1, 1, 1]])

In [54]: result = np.empty_like(idx)

In [55]: for i in range(idx.shape[0]):
    ...:     result[i] = arr[idx[i]]
    ...:

In [56]: result
Out[56]:
array([[6, 9, 9, 7],
    [7, 9, 9, 8],
    [7, 7, 7, 7]])

如何矢量化for 循环？我找不到通过索引矩阵“多次”访问一维数组的方法，其中每行都是一个索引数组。

【问题讨论】：

简单的arr[idx]，输出形状通常由idx形状决定
该死，我不敢相信我什至没有尝试过。

标签： python numpy multidimensional-array vectorization matrix-indexing

【解决方案1】：

如 cmets 中所述，您可以使用 idx 数组简单地索引到数组 arr。

In [47]: arr 
Out[47]: array([6, 7, 8, 9])

In [48]: idx     
Out[48]: 
array([[3, 2, 2, 0],
       [0, 3, 2, 3],
       [3, 2, 2, 3]])

In [49]: arr[idx] 
Out[49]: 
array([[9, 8, 8, 6],
       [6, 9, 8, 9],
       [9, 8, 8, 9]])

如果您想要一种不那么神奇而更有启发性的方法，那么下面的方法会更有帮助。

# flatten the `idx` array; index into `arr`; then reshape to `idx's` shape.

In [50]: arr[idx.ravel()].reshape(idx.shape) 
Out[50]: 
array([[9, 8, 8, 6],
       [6, 9, 8, 9],
       [9, 8, 8, 9]])

【讨论】：