数组太多的指数使用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()部分。