使用索引數組中的索引元組索引多維數組 - NumPy / Python

Question

我有一個尺寸為(6,m,n)的 3-D numpy 數組a 。 我還有一個尺寸為(20,20,20,20,20,20)的 6 (20,20,20,20,20,20) numpy 數組b ，可以有效地用作掩碼。

我想使用第一個數組中每個位置(m,n)的 6 個值來檢索第二個數組中的相應值。 實際上，我會將 3D int 數組壓縮為 2D 布爾數組。 我認為解決方案將使用np.where ，但我認為它不能處理使用值作為索引。

對此的幼稚實現將類似於：

for i in range(m):
    for j in range(n):
         new_arr[i,j]=b[tuple(a[:,i,j])]

有沒有辦法在不使用循環的情況下實現這一點？

Answer 1

方法#1

將a重塑為2D保持第一個軸的長度相同。 將每個這樣的 2D-flattened-block 轉換為一個元組，然后索引到b 。 這種元組轉換導致沿第一個軸打包每個元素作為索引器來選擇每個元素b 。 最后需要重新整形以獲得2D輸出。 因此，實現看起來像這樣 -

b[tuple(a.reshape(6,-1))].reshape(m,n)

或者，跳過所有那些重塑混亂的事情，簡單地做——

b[tuple(a)]

這會創建相同的索引器並解決問題。

方法#2

或者，我們也可以計算扁平索引，然后用這些索引到扁平b並從中提取相關的布爾值 -

b.ravel()[np.ravel_multi_index(a,b.shape)]

大型數據集上的計時 -

In [89]: np.random.seed(0)
    ...: m,n = 500,500
    ...: b = np.random.rand(20,20,20,20,20,20)>0.5
    ...: a = np.random.randint(0,20,(6,m,n))

In [90]: %timeit b[tuple(a)]
14.6 ms ± 184 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [91]: %timeit b.ravel()[np.ravel_multi_index(a,b.shape)]
7.35 ms ± 136 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)