我需要将元素 A 的每一行与元素 B 的所有行相乘。例如,假设我们有以下数组:
A = np.array([[1,5],[3,6]])
B = np.array([[4,2],[8,2]])
我想获得以下数组 C :
C = np.array([[4,10],[8,10],[12,12],[24,12]])
我可以使用 for 循环来做到这一点,但我认为可能会有更好的方法。
编辑:我曾想过重复和平铺,但是我的数组并不小。可能会造成一些内存问题。
答案 0 :(得分:5)
利用broadcasting
将A
的暗度扩展为3D
到None/np.newaxis
,执行逐元素乘法并重整为2D
-
(A[:,None]*B).reshape(-1,B.shape[1])
实际上是-
(A[:,None,:]*B[None,:,:]).reshape(-1,B.shape[1])
示意性地表示为:
A : M x 1 x N
B : 1 x K x N
out : M x K x N
最终重塑形状以合并最后两个轴,并为我们提供(M x K*N)
形的2D
数组。
我们还可以使用einsum
在一个函数调用中执行对3D
的扩展和逐元素乘法-
np.einsum('ij,kj->ikj',A,B).reshape(-1,B.shape[1])