Question

我在NumPy工作。我有一个浮点数组U，形状（n，d，d）和一个形状为（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（bools）的j坐标和U（dxd子阵列）的j坐标。

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

3D Numpy数组到3D Numpy数组

2 个答案: