數組太多的指數使用np.where

Question

時，我有代碼：

a=b=np.arange(9).reshape(3,3) 
c=np.zeros(3) 

for x in range(3): 
    c[x]=np.average(b[np.where(a<x+3)]) 

的c輸出是

>>>array([ 1. , 1.5, 2. ])

取而代之的for循環，我想利用數組（矢量化），然後我做了以下代碼：

a=b=np.arange(9).reshape(3,3) 
c=np.zeros(3) 
i=np.arange(3) 
c[i]=np.average(b[np.where(a<i[:,None,None]+3)]) 

但它顯示了IndexEr ROR：數組

太多的指數作爲a<i[:,None,None]+3

它正確地顯示

array([[[ True, True, True], 
     [False, False, False], 
     [False, False, False]], 

     [[ True, True, True], 
     [ True, False, False], 
     [False, False, False]], 

     [[ True, True, True], 
     [ True, True, False], 
     [False, False, False]]], dtype=bool) 

但是當我使用b[np.where(a<i[:,None,None]+3)]，它再次顯示IndexError：數組太多的索引。我無法獲得c的正確輸出。

Answer 1

我感到你正試圖在這裏向量化東西，雖然沒有明確提到。現在，我認爲你不能以矢量化的方式進行索引。爲了以矢量化的方式解決你的問題，我建議使用matrix-multiplication，並使用np.tensordot，並在broadcasting的幫助下，獲得總和減少的更有效方法，正如您在試驗中已經闡明的那樣。

因此，一種解決方案將是 -

from __future__ import division 

i = np.arange(3) 
mask = a<i[:,None,None]+3 
c = np.tensordot(b,mask,axes=((0,1),(1,2)))/mask.sum((1,2)) 

Related post to understand tensordot。對性能

可能改進

• 轉換掩模供給到np.dot之前浮動D型細胞作爲BLAS基於矩陣的乘法將與它更快。

• 使用np.count_nonzero而不是np.sum計算布爾值。所以，用它來代替mask.sum()部分。