如何用Numpy高效计算两个矩阵所有行对的平方差和？

核心需求 #

需要生成新数组存储两个矩阵所有行两两对应的平方差之和，替代原生双重for循环逻辑，适配行数过万、维度为(20000,1000)级别的大矩阵场景，同时避免内存溢出。

你给出的示例双重for循环逻辑如下：

import numpy as np
a = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]])
b = np.array([[1.2, 2.3, 3.4], [4.5, 5.6, 7.8], [9.10, 10.11, 11.12]])
summed = np.ones((2,3))
for i, aSample in enumerate(a):
    for j, bSample in enumerate(b):
       summed[i, j] = np.sum(np.power(aSample - bSample, 2))

运行得到的预期输出为：

array([[ 18.29  , 112.45  , 308.6765],
       [  7.49  ,  79.65  , 251.0165]])

可用实现方案 #

方案1：广播法（小矩阵可用，大矩阵易内存溢出） #

利用Numpy广播机制直接做差计算，代码简洁，小矩阵下运行速度快，但20000行级别的矩阵运算时会生成(20000,20000,1000)的三维中间数组，内存占用高达数百GB，普通设备无法承载：

summed = np.sum((a[:, np.newaxis] - b)**2, axis=2)

方案2：平方展开法（内存友好，适配大矩阵场景） #

利用平方差公式(a-b)² = a² + b² - 2ab展开计算，通过矩阵乘法和低维广播避免生成高维中间数组，内存占用可控，完全适配你提到的万行级大矩阵场景：

# 计算a每行的平方和，调整形状为(M,1)适配广播
a_sq = np.sum(a**2, axis=1).reshape(-1, 1)
# 计算b每行的平方和，形状为(N,)
b_sq = np.sum(b**2, axis=1)
# 计算a和b的行点积矩阵，形状为(M,N)
ab_dot = a @ b.T
# 代入公式直接得到最终结果
summed = a_sq + b_sq - 2 * ab_dot

该方案运算效率和内存占用均为工业界计算欧氏距离矩阵的标准实现，输出结果和双重for循环完全一致。

内容的提问来源于stack exchange，提问作者Toothery