基于两个CSV计算最小欧氏距离的高效Python方案求助

问题背景 #

我有两个CSV文件：

CSV A #

Col1   Col2         Col3        Col4         Col5        Col6
23      XY      -2852059.23, 5257582.12, -2205236.45      AA
25      AB      -2852159.23, 5257577.12, -2205245.43      BB
31      CF      -2852259.23, 5257569.12, -2205240.46      CC 
35      CV      -2852359.23, 5257565.12, -2205246.47      DD
76      FH      -2852459.23, 5257573.12, -2205241.47      EE
34      TY      -2852559.23, 5257564.12, -2205244.56      FF

CSV B #

Col1       Col2      Col3      Col4
75      -28059.23, 52582.12, -25236.45
23      -28159.23, 52577.12, -25245.43
34      -28259.23, 52569.12, -25240.46
56      -28359.23, 52565.12, -25246.47
67      -28459.23, 52573.12, -25241.47
34      -28559.23, 52564.12, -25244.56

需求：使用欧氏距离公式 $\sqrt{(x₂ - x₁)² + (y₂ - y₁)² + (z₂ - z₁)²}$，计算CSV A每行与CSV B所有行的距离，然后取每行的最小距离。

已尝试方法及问题 #

1. 逐行遍历计算

df1['min_sqrt'] = np.sqrt((df2['Col2'].apply(lambda x: (x-df1['Col3'])**2))+(df2['Col3'].apply(lambda x: (x-df1['Col4'])**2))+(df2['Col4'].apply(lambda x: (x-df1['Col5'])**2))).min()

功能正常，但处理大数据时运行速度极慢。

2. Cross Merge + 分组聚合

out = (
    df1.merge(df2, how="cross")
       .assign(x= lambda df: (df['Col2'] - df['Col3']).pow(2),
               y= lambda df: (df['Col3'] - df['Col4']).pow(2),
               z= lambda df: (df['Col4'] - df['Col5']).pow(2),
               distance= lambda df: np.sqrt(df.x + df.y + df.z))
       .groupby(["Col1", "Col2", "Col3", "Col4", "Col5", "Col6"], sort=False, axis=0).agg(min_sqrt=("distance", "min"))
                
   )

小数据集可正常运行，但大数据下触发内存错误：

MemoryError: Unable to allocate 23.2 GiB for an array with shape (31, 100453255) and data type float64

高效解决方案 #

方案1：NumPy向量化运算 #

利用NumPy的广播机制实现无循环计算，内存效率远高于Cross Merge：

import numpy as np
import pandas as pd

# 读取CSV并提取三维坐标列（需确保坐标列是数值类型）
coords_a = df1[['Col3', 'Col4', 'Col5']].to_numpy()
coords_b = df2[['Col2', 'Col3', 'Col4']].to_numpy()

# 计算距离平方矩阵：(n,3) - (m,3) → (n,m,3)，求和后得到(n,m)的距离平方矩阵
dist_sq = np.sum((coords_a[:, np.newaxis] - coords_b)**2, axis=2)
# 提取每行最小距离并开根号
min_dist = np.sqrt(dist_sq.min(axis=1))

# 将结果合并到原DataFrame
df1['min_sqrt'] = min_dist

优势：纯向量化运算速度快；内存占用为O(n*m)，但NumPy数组的内存效率远高于Pandas DataFrame，适合中等规模数据集。

方案2：KDTree最近邻搜索（超大数据集首选） #

使用Scipy的KDTree进行高效最近邻查询，时间复杂度为O(n log m)，远低于暴力计算的O(n*m)：

from scipy.spatial import KDTree
import pandas as pd

# 读取CSV并提取三维坐标列
coords_a = df1[['Col3', 'Col4', 'Col5']].to_numpy()
coords_b = df2[['Col2', 'Col3', 'Col4']].to_numpy()

# 构建KDTree结构
tree = KDTree(coords_b)
# 查询每个点的最近邻距离（k=1表示仅取最近的1个点）
min_dist, _ = tree.query(coords_a, k=1)

# 合并结果到原DataFrame
df1['min_sqrt'] = min_dist

优势：内存占用低，仅需存储KDTree结构；查询速度随数据量增大优势越明显，适合百万级以上的超大数据集。

预处理注意事项 #

如果CSV中的坐标是在单个列中用逗号分隔的字符串（如示例中CSV A的Col3），需先拆分并转换为数值类型：

# 处理CSV A的坐标列
df1[['x', 'y', 'z']] = df1['Col3'].str.split(',', expand=True).astype(float)
coords_a = df1[['x', 'y', 'z']].to_numpy()

# 处理CSV B的坐标列
df2[['x', 'y', 'z']] = df2['Col2'].str.split(',', expand=True).astype(float)
coords_b = df2[['x', 'y', 'z']].to_numpy()

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