自定义点积的Numpy矩阵乘法

默认的矩阵乘法计算为

c[i,j] = sum(a[i,k] * b[k,j])

我正在尝试使用自定义公式而不是点积来获得

c[i,j] = sum(a[i,k] == b[k,j])

---

在numpy有没有一种有效的方法可以做到这一点？

您可以使用广播：

c = sum(a[...,np.newaxis]*b[np.newaxis,...],axis=1)  # == np.dot(a,b)

c = sum(a[...,np.newaxis]==b[np.newaxis,...],axis=1)

---

我在

b

中包含了

newaxis

，只需说明数组是如何扩展的。还有其他向数组添加维度的方法（重塑、重复等），但效果相同。将

a

和

b

展开到相同的形状，进行逐元素相乘（或==），然后在正确的轴上求和

谢谢你。这个很整洁。