索引二維數組由索引的二維數組沒有循環

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_的文檔以瞭解這一點。