由没有循环的2D索引数组索引2D numpy数组

问题描述：

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

I am looking for a vectorized way to index a numpy.array by numpy.array of indices.

例如:

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)的行索引.我想要的输出是:

I want to build a new array, such that every row(i) in that array is a row(i) of array a, indexed by row of array inds(i). My desired output is:

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

我可以通过循环来实现:

I can achieve this with a loop:

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

但是我正在寻找一种纯粹的矢量化解决方案.

But I am looking for a pure vectorized solution.

答

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

When using arrays of indices to index another array, the shape of each index array should match the shape of the output array. You want the column indices to match inds, and you want the row indices to match the row of the output, something like:

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

由于广播，您只能使用上面的一列，所以您可以将np.arange(3)[:,None]用作垂直的arange，因为None会插入新轴:

You can just use a single column of the above, due to broadcasting, so you can use np.arange(3)[:,None] is the vertical arange because None inserts a new axis:

>>> 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]]

由没有循环的2D索引数组索引2D numpy数组

相关推荐