索引二维数组由索引的二维数组没有循环

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的行（ⅰ索引）。我期望的输出是：

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_的文档以了解这一点。