通過沒有循環的二維索引數組索引二維 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 ，通過陣列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_文檔以了解這一點。