Умножьте 2dim-массив с вектором, чтобы получить массив 3dim

Question

Предположим, у меня есть массив numpy, который равен 2d, и массив 1d

In [127]: A = np.array([[1, 2],[3, 4]])

In [128]: B = np.array([10, 100])

Чего я хотел бы добиться, так это получить трехмерный массив C, где C[:, :, 0] = A*B[0] и C[:, :, 1] = A*B[1]. Я смог сделать это через np.einsum, но выглядит как перебор.

In [129]: np.einsum('ij, k -> ijk', A, B)[:, :, 0]
Out[129]: 
array([[10, 20],
       [30, 40]])

In [130]: np.einsum('ij, k -> ijk', A, B)[:, :, 1]
Out[130]: 
array([[100, 200],
       [300, 400]])

Есть ли более простая версия?

Answer 1

Чтобы использовать B для масштабирования по первой оси, мы можем просто использовать broadcasting примерно так -

B[:,None,None]*A # with einsum : np.einsum('ij, k -> kij', A, B)

Чтобы получить эквивалент C[:, :, 0] = A*B[0] and C[:, :, 1] = A*B[1], нам нужно расширить A вместо -

A[...,None]*B # with einsum : np.einsum('ij, k -> ijk', A, B)

Если бы здесь не происходило уменьшения суммы, основанные на broadcasting были бы быстрее, чем einsum.