沿B广播A的每一行，避免'重复'

Question

我有2个2D阵列，1个相同维度的轴：

a = np.array(np.arange(6).reshape((2,3)))
b = np.array(np.arange(12).reshape((3,4)))

我希望将a的每一行与b相乘并广播，即

b_r = np.repeat(b[:,:,None], 2, axis=2)
ab = a.T[:,None,:] * b_r

是否可以在避开repeat的同时进行广播？我们的想法是避免repeat操作的不必要的内存分配。

Answer 1

您可以直接输入b[:,:,None]而不重复，因为broadcasting的定义会为您播放。

因此，只需做 -

ab = a.T[:,None,:]*b[:,:,None]

我们可以通过跳过:的尾随a并使用...将:,:替换为b来实现一点紧凑，就像这样 -

ab = a.T[:,None]*b[...,None]

对于 kicks ，这里有一个使用np.einsum，这个性能稍差，但是一旦我们超过它的字符串表示法就会更具表现力 -

ab = np.einsum('ij,jk->jki',a,b)