python - 通过没有循环的二维索引数组索引二维 numpy 数组

Question

我正在寻找一种矢量化的方式来通过索引的 numpy.array 对 numpy.array 进行索引。

例如:

import numpy as np

a = np.array([[0,3,4],
              [5,6,0],
              [0,1,9]])

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

我想构建一个新数组，这样该数组中的每一行 (i) 都是数组 a 的一行 (i)，由数组 inds(i) 的行索引。我想要的输出是:

array([[ 0.,  3.],   # a[0][:,[0,1]]
       [ 6.,  0.],   # a[1][:,[1,2]] 
       [ 0.,  9.]])  # a[2][:,[0,2]]

我可以用一个循环来实现:

def loop_way(my_array, my_indices):
    new_array = np.empty(my_indices.shape)
    for i in xrange(len(my_indices)):
        new_array[i, :] = my_array[i][:, my_indices[i]]
    return new_array 

但我正在寻找一个纯矢量化的解决方案。

Answer 1

当使用索引数组索引另一个数组时，每个索引数组的形状应与输出 数组的形状相匹配。您希望列索引与 inds 匹配，并且您希望行索引与输出行匹配，例如:

array([[0, 0],
       [1, 1],
       [2, 2]])

你可以只使用上面的单个列，由于广播，所以你可以使用 np.arange(3)[:,None] 是垂直 arange 因为 None 插入了一个新轴:

>>> np.arange(3)[:, None]
array([[0],
       [1],
       [2]])

最后，一起:

>>> a[np.arange(3)[:,None], inds]
array([[0, 3],   # a[0,[0,1]]
       [6, 0],   # a[1,[1,2]] 
       [0, 9]])  # a[2,[0,2]]