没有for循环的numpy行对平方行差异总和（仅api调用）

Question

对于那些可以阅读 Latex 的人，这是我想要计算的：

$$k_{xyi} = \sum_{j}\left ( \left ( x_{i}-x_{j} \right )^{2}+\left ( y_{i}-y_{j} \right )^{2} \right )$$

其中 x 和 y 是矩阵 A 的行。

对于只有计算机语言的人，这将翻译为： k(x,y,i) = sum_j( (xi - xj)^2 + (yi - yj)^2 ) 其中 x 和 y 是矩阵 A 的行。

所以 k 是一个 3d 矩阵。

这只能通过 API 调用来完成吗？ （没有 for 循环）

这里是测试启动：

import numpy as np
A = np.random.rand(4,4)
k = np.empty((4,4,4))
for ix in range(4):
    for iy in range(4):
        x = A[ix,]
        y = A[iy,]
        sx = np.power(x - x[:,np.newaxis],2)
        sy = np.power(y - y[:,np.newaxis],2)
        k[ix,iy] = (sx + sy).sum(axis=1).T

现在对于编码大师，请将两个 for 循环替换为 numpy API 调用。

更新： 忘了说我需要一种节省 RAM 空间的方法，我的 A 矩阵通常是 20-30 千平方。因此，如果您的答案不会创建巨大的临时多维数组，那就太好了。

Answer 1

我会改变你的乳胶看起来更像下面的东西 - 它更容易混淆 imo：

据此，我假设您表达式中的最后一行应该是：

k[ix,iy] = (sx + sy).sum(axis=-1)

如果是这样，您可以按如下方式计算上述表达式：

Axij = (A[:, None, :] - A[..., None])**2
k = np.sum(Axij[:, None, :, :] + Axij, axis=-1)

上面首先扩展了一个内存密集型 4D 数组。如果您担心内存问题，可以通过引入新的 for 循环跳过此步骤：

k = np.empty((4,4,4))
Axij = (A[:, None, :] - A[..., None])**2
for xi in range(A.shape[0]):
    k[xi] = np.sum(Axij[xi, None, :, :] + Axij, axis=-1)

这会更慢，但不会像你想象的那么快，因为你仍然在 numpy.您可能可以跳过 3D Axij 中间，但这样做会再次降低性能。

如果您的矩阵在边缘上真的是 20k，那么您的 3D 输出将是 64TB。您不会在 numpy 甚至内存中执行此操作（除非您有大规模的分布式内存系统）。