python - 对平方数组求和

Question

作为批量欧氏距离计算的一部分，我正在计算

(X * X).sum(axis=1)

其中 X 是一个相当大的二维数组。这工作正常，但它构造了一个与 X 大小相同的临时数组。有什么方法可以摆脱这种临时性，但又能保持矢量化操作的效率？

明显的候选人，

np.array([np.dot(row, row) for row in X])

有效，但使用 Python 列表作为临时列表，使其相当慢。

如果没有轴，内存效率的形式将是

(X * X).sum()  =>  np.dot(X.ravel(), X.ravel())

我知道，当 axis=1 时，它等同于

np.diag(np.dot(X, X.T))

这让我研究了 dot 的概括，例如 np.inner、np.tensordot 和 np.einsum，但我不知道他们将如何解决我的问题。

Answer 1

einsum 等价于:

np.einsum('ij,ij->i', X, X)

虽然我不确定这在内部是如何工作的，所以它可能会或可能不会解决您的问题。