numpy - 带有自定义点积的 Numpy 矩阵乘法

Question

默认矩阵乘法计算为

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

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

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

在 numpy 中有没有一种有效的方法来做到这一点？

Answer 1

您可以使用广播：

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

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

我包括在内newaxis只是b为了清楚地说明该数组是如何扩展的。还有其他向数组添加维度的方法（重塑、重复等），但效果是一样的。展开a并b以相同的形状进行逐个元素的乘法（或 ==），然后在正确的轴上求和。