Numpy查找2D三维点集公共点索引：兼容浮点误差的向量化实现

系统环境 #

• 操作系统: Windows 10 (x64), Build 1909
• Python 版本: 3.8.10
• Numpy 版本: 1.21.2

问题描述 #

现有两个形状为(N, 3)的Numpy二维浮点数组，存储(x, y, z)格式的三维坐标点，需要找到大数组中与小数组内点匹配的索引。要求为Pythonic的向量化实现，适配真实数据集存在的浮点误差场景，同时性能可支撑百万级点集的计算。

当前遇到的问题：直接使用np.isin无法识别存在微小浮点误差的匹配点，直接调用np.isclose会触发形状广播错误，现有列表推导式方案性能不足以支撑大规模数据集。

可行实现方案 #

以下两种方案均支持浮点误差容差配置，可根据数据规模选择：

方案1：纯Numpy广播实现（适用十万级及以下数据集）

通过维度扩展触发广播计算，无需额外依赖，实现逻辑简单：

import numpy as np

def match_points_broadcast(large_arr: np.ndarray, small_arr: np.ndarray, atol: float = 1e-6) -> tuple[np.ndarray, np.ndarray]:
    """
    匹配两个三维坐标数组的点
    :param large_arr: 大数组，形状(M, 3)
    :param small_arr: 小数组，形状(N, 3)
    :param atol: 绝对容差，适配浮点误差
    :return: 大数组中匹配成功的索引，对应小数组的匹配位置索引
    """
    # 维度扩展后计算每个坐标的接近程度，三个坐标均接近则判定为匹配
    close_mask = np.isclose(large_arr[:, None, :], small_arr[None, :, :], atol=atol).all(axis=-1)
    # 提取大数组匹配索引
    large_match_idx = np.where(close_mask.any(axis=1))[0]
    # 提取对应小数组的匹配索引
    _, small_match_idx = np.where(close_mask)
    return large_match_idx, small_match_idx

注意：该方案会生成形状为(M, N, 3)的临时数组，当数据量达到百万级时内存占用会非常高，建议仅在中小规模场景下使用。

方案2：cKDTree实现（适用百万级大规模数据集）

基于scipy.spatial.cKDTree的空间索引查询，时间复杂度为O(M log N)，内存占用极低，是大规模场景的最优选择：

import numpy as np
from scipy.spatial import cKDTree

def match_points_kdtree(large_arr: np.ndarray, small_arr: np.ndarray, atol: float = 1e-6) -> tuple[np.ndarray, np.ndarray]:
    """
    匹配两个三维坐标数组的点
    :param large_arr: 大数组，形状(M, 3)
    :param small_arr: 小数组，形状(N, 3)
    :param atol: 绝对容差，适配浮点误差
    :return: 大数组中匹配成功的索引，对应小数组的匹配位置索引
    """
    # 基于小数组构建KD树
    kd_tree = cKDTree(small_arr)
    # 查询大数组每个点在容差范围内的最近点
    dist, small_match_idx = kd_tree.query(large_arr, distance_upper_bound=atol)
    # 过滤匹配成功的结果
    match_mask = dist <= atol
    large_match_idx = np.where(match_mask)[0]
    return large_match_idx, small_match_idx[match_mask]

该方案针对百万级点集可在数秒内完成计算，性能远高于循环/列表推导实现。

内容的提问来源于stack exchange，提问作者adam.hendry