大数据量下Python DataFrames行对差值（欧氏距离）计算优化咨询

兄弟，你这情况我太懂了——3万行×3万行的行对，要是直接生成三维差值数组（30000×30000×8），光是内存就得占50多GB，别说计算了，内存都扛不住。而且循环遍历每一对行更是慢到离谱，得换思路，用向量化的线性代数运算，把复杂度和内存占用都降下来。

核心思路：用欧氏距离的数学公式绕开三维数组 #

欧氏距离的平方可以拆解成这样的公式：

||a - b||² = ||a||² + ||b||² - 2×(a·b)

其中：

• ||a||²是a行元素的平方和（L2范数的平方）
• a·b是a行和b行的点积

这个公式的好处是，我们可以用矩阵运算一次性算出所有行对的结果，不需要生成中间的三维差值数组，内存占用直接从O(n²×d)降到O(n²)（n是行数，d是列数），计算速度也会因为numpy的BLAS优化（底层是C/汇编实现）大幅提升。

具体实现方法 #

首先把DataFrame转成numpy数组（pandas的运算在这种大规模场景下不如numpy直接）：

import numpy as np
import pandas as pd

# 假设你的两个DataFrame是df1和df2
A = df1.to_numpy()
B = df2.to_numpy()

方法1：手动用矩阵运算实现

# 计算每行的平方和（L2范数平方）
A_norm = np.sum(A ** 2, axis=1)
B_norm = np.sum(B ** 2, axis=1)

# 计算所有行对的点积矩阵（A的每行和B的每行做点积）
dot_product = A @ B.T

# 套用公式计算欧氏距离的平方
dist_sq = A_norm[:, np.newaxis] + B_norm[np.newaxis, :] - 2 * dot_product

# 如果需要实际的欧氏距离，再开根号（开根号是可选的，比如排序时用平方距离更高效）
dist = np.sqrt(dist_sq)

方法2：用scipy的现成工具（更简洁）

scipy的cdist函数专门做两两距离计算，内部已经做了极致优化，甚至支持多线程，代码更省心：

from scipy.spatial.distance import cdist

# 直接计算所有行对的欧氏距离
dist = cdist(A, B, metric='euclidean')

内存优化技巧 #

如果3万×3万的距离矩阵（9亿个元素）还是超出你的内存（大概7GB左右，按float64算），可以分块计算：

• 把A分成若干小块（比如每次处理1000行），每块和B计算距离，然后把结果拼接起来
• 这样每次只需要处理1000×30000的矩阵，内存占用瞬间降到2.4GB左右

示例代码：

block_size = 1000
dist_blocks = []

for i in range(0, A.shape[0], block_size):
    A_block = A[i:i+block_size]
    block_dist = cdist(A_block, B, metric='euclidean')
    dist_blocks.append(block_dist)

# 拼接所有块得到完整距离矩阵
dist = np.vstack(dist_blocks)

额外加速方案：GPU加速 #

如果你的机器有NVIDIA GPU，可以用cupy代替numpy，矩阵运算速度能提升几十倍——只需要把数组转成cupy数组，代码几乎不用改：

import cupy as cp

A_cp = cp.array(A)
B_cp = cp.array(B)

A_norm_cp = cp.sum(A_cp ** 2, axis=1)
B_norm_cp = cp.sum(B_cp ** 2, axis=1)
dot_product_cp = A_cp @ B_cp.T
dist_sq_cp = A_norm_cp[:, None] + B_norm_cp[None, :] - 2 * dot_product_cp
dist_cp = cp.sqrt(dist_sq_cp)

# 转回numpy数组（如果需要的话）
dist = cp.asnumpy(dist_cp)

注意事项 #

• 如果你的DataFrame里有缺失值，一定要先处理（比如df.fillna(0)或者用均值填充），否则计算结果会出现NaN
• 优先用float32类型（如果精度允许），内存占用直接减半，计算速度也会更快：A = df1.to_numpy(dtype=np.float32)

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