使用布尔掩码进行矢量化reduce乘积的最快方法

Question

我有一个3Dnumpy数组A和2Dnumpy布尔掩码B。

A的前两个维度与B匹配

我想知道对于A的第一个维度，有没有什么快速的方法，根据B沿着第二个维度选择第三个维度，在第二个维度上执行一个简化的乘积。

我期望的输出C是一个二维numpy数组，第一个维度是A，第二个维度是A的第三个维度。

我目前的解决方案是C = np.prod(A*np.repeat(B[...,np.newaxis], A.shape[-1], 2), 1)

有没有更好的选择？

Answer 1

举个具体例子：

In [364]: A=np.arange(1,25).reshape(2,3,4); B=np.arange(1,7).reshape(2,3)
In [365]: C = np.prod(A*np.repeat(B[...,np.newaxis], A.shape[-1], 2), 1)

这个重复做了：

In [366]: np.repeat(B[...,np.newaxis], A.shape[-1], 2)
Out[366]: 
array([[[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3]],

       [[4, 4, 4, 4],
        [5, 5, 5, 5],
        [6, 6, 6, 6]]])
In [367]: _.shape
Out[367]: (2, 3, 4)
In [368]: A*np.repeat(B[...,np.newaxis], A.shape[-1], 2)
Out[368]: 
array([[[  1,   2,   3,   4],
        [ 10,  12,  14,  16],
        [ 27,  30,  33,  36]],

       [[ 52,  56,  60,  64],
        [ 85,  90,  95, 100],
        [126, 132, 138, 144]]])

但是根据broadcasting规则，repeat是不需要的：

In [369]: A*B[...,np.newaxis]
Out[369]: 
array([[[  1,   2,   3,   4],
        [ 10,  12,  14,  16],
        [ 27,  30,  33,  36]],

       [[ 52,  56,  60,  64],
        [ 85,  90,  95, 100],
        [126, 132, 138, 144]]])
In [371]: np.prod(_369, axis=1)
Out[371]: 
array([[   270,    720,   1386,   2304],
       [556920, 665280, 786600, 921600]])

您可以将prod分别应用于A和B，但我不知道这是否会有很大区别：

In [373]: np.prod(A,1)*np.prod(B,1)[:,None]
Out[373]: 
array([[   270,    720,   1386,   2304],
       [556920, 665280, 786600, 921600]])