通过没有循环的2D索引数组索引2D 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)都是数组的行(i)a,由数组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]]

Answer 2

这样做是可能的，虽然有点不明显，如下所示：

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

索引np.arange(a.shape[0])只是索引列索引数组inds适用的行。追加[:, None]修改此数组的形状，使其形状为(a.shape[0], 1)，即每个行索引位于 1 列宽 2D 数组的单独行中。

基本原则是索引数组的维数必须一致，它们的形状也必须一致。请参阅文档np.ix_以了解这一点。