Numpy matrix multiplication with custom dot product

Question

The default matrix multiplication is computed as

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

I am trying to use a custom formula instead of the dot product to get

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

Is there an efficient way to do this in numpy?

Answer 1

You could use broadcasting:

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

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

I included the newaxis in b just make it clear how that array is expanded. There are other ways of adding dimensions to arrays (reshape, repeat, etc), but the effect is the same. Expand a and b to the same shape to do element by element multiplying (or ==), and then sum on the correct axis.