python - 编写以下内容的最快方法

我在python中执行以下操作：

tmp = np.empty_like(J,dtype=X.dtype)
for idx, (ii, jj) in enumerate(zip(I, J)):
    tmp[idx] = sum((X[ii] - X[jj])**2)

其中X是50000 x 128 numpy数组
和I和J是大小为（763690）的整数numpy数组（稀疏矩阵的列和行）

现在的问题是，上述操作大约需要30秒钟才能完成，而且我看不到需要做些什么来加快速度。我知道它可以更快地完成，因为我在Matlab中有一个类似的代码，几乎不需要任何时间。

我在这里做错了什么？
是否与内存步幅访问有关？是否不使用内置函数？或者是其他东西？我应该对其进行并行化/矢量化吗？

（我知道标题很糟糕，但是我想不出一种好方法写，欢迎提出建议！）

最佳答案

我们可以这样做：

np.sum((X[I]-X[J])**2, axis=1)

因此，这里我们首先使用下标生成763 690×128矩阵X[I]，其中对于I中的每个项目，我们都使用X中的相应行。 X[j]也会发生同样的情况。然后，我们将两者相减，得到763 690×128矩阵。我们可以按元素对这些项进行平方，然后在第一个轴上计算总和。因此，这意味着对于每一行，我们都获得一个值。因此，结果是具有763 690个元素的平面阵列。

关于python - 编写以下内容的最快方法，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/57931055/