基于近似坐标的Pandas DataFrame数据掩码方法求助

问题背景 #

有两个存储XYZ坐标的Pandas DataFrame（df1和df2），需要用df2中的坐标对df1进行掩码移除操作，但两组坐标存在微小数值差异，直接取整匹配容易出现失败情况。示例数据如下：

import pandas as pd
import numpy as np

df1 = pd.DataFrame(data=None, columns=['x', 'y', 'z'])
df1.x = [104245, 252355, 547364, 135152]
df1.y = [842714, 135812, 425328, 124912]
df1.z = [125125, 547574, 364343, 346372]

df2 = pd.DataFrame(data=None, columns=['x', 'y', 'z'])
df2.x = [104230, 547298]
df2.y = [842498, 424989]
df2.z = [124976, 364001]

原方案通过对坐标取整后匹配，但常因取整粒度问题导致匹配失败，需要基于最近邻距离的更可靠掩码方法。

解决方案：KDTree最近邻匹配 #

利用scipy.spatial.KDTree实现高效的3维坐标最近邻搜索，通过设定距离阈值来判定是否为需要移除的匹配点，步骤如下：

1. 导入依赖库 #

from scipy.spatial import KDTree

2. 提取坐标数组并构建KDTree #

将df2的坐标转换为numpy数组，构建KDTree用于快速查询：

# 提取df1和df2的XYZ坐标数组
coords1 = df1[['x', 'y', 'z']].values
coords2 = df2[['x', 'y', 'z']].values

# 用df2的坐标构建KDTree
tree = KDTree(coords2)

3. 查询最近邻距离并生成掩码 #

对df1的每个坐标点，查询其在df2中的最近邻距离，设定合理的距离阈值（根据数据的误差范围调整，示例中设为500），距离小于阈值的点标记为需要移除：

# 查询每个df1点到df2的最近邻距离和对应索引
distances, indices = tree.query(coords1, k=1)

# 设定距离阈值，这里根据示例数据的差异设为500
distance_threshold = 500

# 生成掩码：保留距离大于阈值的点（即未匹配到df2的点）
mask = distances > distance_threshold

# 应用掩码得到最终结果
df_filtered = df1[mask]

关键说明 #

• 距离阈值的选择：需要根据实际数据的误差范围调整，比如如果坐标误差在±100以内，阈值可设为200，确保覆盖合理的差异范围。
• 性能优化：KDTree的查询时间复杂度为O(n log n)，适合处理大规模坐标数据，比暴力匹配效率高得多。
• 多匹配处理：如果需要确保df2的每个点只匹配一次df1中的点，可以使用tree.query_ball_point结合去重逻辑，避免重复移除。

内容的提问来源于stack exchange，提问作者Marcus K.