使用 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)])

但它显示 IndexError: too many indices for array

至于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: too many indices for array。我无法得到c 的正确输出。

Answer 1

我感觉到您正在尝试将这里的东西矢量化，尽管没有明确提及。现在，我认为您不能以矢量化方式进行索引。为了以矢量化的方式解决您的问题，我建议使用np.tensordot 并在broadcasting 的帮助下使用matrix-multiplication 进行减和，正如您在试验中已经提出的那样。

因此，一种解决方案是 -

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 之前将掩码转换为 float dtype，因为使用它基于 BLAS 的矩阵乘法会更快。

• 使用np.count_nonzero 而不是np.sum 来计算布尔值。所以，用它来替换mask.sum() 部分。