如何对这个python代码进行矢量化以使其更高效？(速度快)

Question

我试图使用numpy模块向量化我的代码,原始代码如下:

m = [M[i,:9].dot(N[:9,i]) for i in xrange(9)]

我改进了代码:

m = np.diagonal(M[:9,:9].dot(N[:9,:9]))  

然而,这将导致一些不必要的计算(特别是当指数远大于9时).我该怎么做才能进一步提高效率？

编辑:基本上我打算做的是计算两个矩阵M和N的点积的对角元素.

Answer 1

您可以np.einsum根据需要使用第一轴M与第二轴对齐N,同时减少/丢失输入的剩余轴.因此,我们将有一个einsum基础的解决方案,如下所示:

m = np.einsum('ij,ji->i',M[:,:9],N[:9,:])