3D Numpy阵列到3D Numpy阵列

Question

我在NumPy工作。 我有一个形状为（n，d，d）的浮点数U数组，以及一个形状为（k，n）的2D布尔数组B. 看起来有点像

U = np.array([
    [[0,1],
    [2,3]
    ],
    [[4,5],
    [6,7]
    ]
    [[1,2],
    [3,4]
    ]
])

B = np.array([
    [True,False,False],
    [True,False,True],
    [True,True,False],
    [False,False,True]
])

我想要一个向量化函数vector_sum（A，B），它将输出形状为（4,2,2）的数组Z，其中Z [0]为U [0]； Z [1]为U [0] + U [2]； Z [2]为U [0] + U [1]，Z [3]为U [2]。 我怎样才能做到这一点？ 我猜想np.einsum有办法做到这一点，但是我不太了解它是如何工作的，所以我有点时间紧迫。

谢谢！

Answer 1

IIUC，您绝对可以使用np.einsum ：

In [70]: np.einsum('ij,jkl->ikl', B, U)
Out[70]: 
array([[[ 0,  1],
        [ 2,  3]],

       [[ 1,  3],
        [ 5,  7]],

       [[ 4,  6],
        [ 8, 10]],

       [[ 1,  2],
        [ 3,  4]]])

它将作用于B的j坐标（布尔）和U的j坐标（dxd子数组）。

Answer 2

这样做：

import numpy as np

U = np.array([
    [[0,1],[2,3]],
    [[4,5],[6,7]],
    [[1,2],[3,4]]
    ])

B = np.array([
    [True,False,False],
    [True,False,True],
    [True,True,False],
    [False,False,True]
])


Z = np.array([U[i].sum(axis=0) for i in B])