使用行索引的2D陣列numpy的高級索引而無需廣播輸出

Question

我有NDIM 3的ndarray array，及指標ndarray idxs與NDIM 2，這對於array第一維度指定索引。第一維idxs與第二維array相匹配，即idxs.shape[0] == array.shape[1]。

我希望得到一個導致ndarray result與NDIM 3和形狀(idxs.shape[1], array.shape[1], array.shape[2])這樣的：

for i0 in range(idxs.shape[1]): 
    for i1 in range(array.shape[1]): 
     result[i0, i1] = array[idxs[i1, i0], i1] 

我怎樣才能更直接地得到這個？

我想過使用高級索引，但我不完全確定它會是什麼樣子。

在Theano，以下工作：

dim1 = theano.tensor.arange(array.shape[1]) 
result = array[idxs[dim1], dim1] 

Answer 1

你for循環做到這一點：

out[i, j] == array[idxs[j, i], j] 

也就是說，在J，I ^個在idxs元素給出行 index爲array爲i，j ^th元素爲out。相應的列列索引到array只是0和idxs.shape[0] - 1之間的序列整數（在這種情況下恰好與array.shape[1] - 1相同，但不一定總是）。

你for環因此像這樣一個數組索引操作來代替：

def simplified(array, idxs): 
    return array[idxs.T, np.arange(idxs.shape[0])] 

我們可以測試的正確性和速度的對抗@ Divakar的回答的功能：

m, n = 500, 400 
array = np.random.rand(m, n) 
idxs = np.random.randint(n, size=(n, m)) 

print(np.allclose(forloop(array, idxs), simplified(array, idxs))) 
# True 

%timeit forloop(array, idxs) 
# 10 loops, best of 3: 101 ms per loop 

%timeit broadcasted_indexing(array, idxs) 
# 100 loops, best of 3: 4.1 ms per loop 

%timeit simplified(array, idxs) 
# 1000 loops, best of 3: 1.66 ms per loop 

Answer 2

創建對應於行索引指數的2D網格：idxs[i1, i0]並使用N x 1陣列列索引。當像這樣索引到array時，列索引將是broadcasted到行索引的形狀。因此，我們將有一個broadcasted indexing爲基礎的方法，像這樣 -

# Get 2D grid of row indices corresponding to two nested loops 
row_idx = idxs[np.arange(array.shape[1])[:,None],np.arange(idxs.shape[1])] 

# Use column indices alongwith row_idx to index into array. 
# The column indices would be broadcasted when put as Nx1 array. 
result = array[row_idx,np.arange(array.shape[1])[:,None]].T 

請注意，在通過@ali_m的評論中提到，np.ix_還可以用來創建row_idx，像這樣 -

row_idx = idxs[np.ix_(np.arange(array.shape[1]),np.arange(idxs.shape[1]))] 

運行測試和驗證輸出

功能定義：

def broadcasted_indexing(array,idxs): 
    row_idx = idxs[np.arange(array.shape[1])[:,None],np.arange(idxs.shape[1])] 
    return array[row_idx,np.arange(array.shape[1])[:,None]].T 

def forloop(array,idxs): 
    result = np.zeros((idxs.shape[1],array.shape[1])) 
    for i0 in range(idxs.shape[1]): 
     for i1 in range(array.shape[1]): 
      result[i0, i1] = array[idxs[i1, i0], i1] 
    return result 

運行測試和驗證輸出：

In [149]: # Inputs 
    ...: m = 500 
    ...: n = 400 
    ...: array = np.random.rand(m,n) 
    ...: idxs = np.random.randint(0,array.shape[1],(n,m)) 
    ...: 

In [150]: np.allclose(broadcasted_indexing(array,idxs),forloop(array,idxs)) 
Out[150]: True 

In [151]: %timeit forloop(array,idxs) 
10 loops, best of 3: 136 ms per loop 

In [152]: %timeit broadcasted_indexing(array,idxs) 
100 loops, best of 3: 5.01 ms per loop