本文介绍如何将 scikit-learn KDTree.query_radius() 返回的嵌套数组（含空数组）一键展平、合并并去重，避免显式 Python 循环，提升代码简洁性与运行效率。
本文介绍如何将 scikit-learn `kdtree.query_radius()` 返回的嵌套数组（含空数组）一键展平、合并并去重，避免显式 python 循环，提升代码简洁性与运行效率。
在使用 sklearn.neighbors.KDTree 进行范围查询时，query_radius() 方法会为每个查询点 Y[i] 返回一个一维 NumPy 数组（dtype=int64），表示所有落在半径 r 内的参考点索引。由于不同查询点的邻域大小不同，该方法返回的是 object 类型的 NumPy 数组（即 np.ndarray of np.ndarray），其中部分元素可能是空数组（如 array([], dtype=int64)），这给后续批量处理（如统计、去重、索引映射）带来不便。
直接对 idxs 调用 np.concatenate(idxs) 会报错（因 concatenate 不支持 object 数组中混有空数组），而手动遍历 for arr in idxs: if len(arr) > 0: ... 虽可行但冗余且低效。推荐方案是使用 numpy.hstack() —— 它专为拼接一维数组设计，能自动跳过空数组，且对 object 数组中的非空子数组进行无缝水平堆叠。
✅ 正确做法如下：
import numpy as np from sklearn.neighbors import KDTree # 示例数据（2D 坐标） X = np.random.rand(49, 2) # 参考点集 Y = np.random.rand(50, 2) # 查询点集 kd_tree = KDTree(X) idxs = kd_tree.query_radius(Y, r=0.2) # 返回 shape=(len(Y),) 的 object 数组 # 一步完成：展平 + 合并 + 去重 unique_neighbors = np.unique(np.hstack(idxs)) print(unique_neighbors) # 输出示例：[ 0 1 2 ... 47] （长度 ≤ len(X)，无重复、无空值）
⚠️ 注意事项：
- np.hstack(idxs) 要求 idxs 中每个元素均为一维数组（query_radius() 默认满足），且自动忽略空数组，安全可靠；
- 若需保留原始邻域结构（如计算每个 Y[i] 的邻居数量），请勿提前展平；本方案仅适用于“全局聚合”场景（如获取所有被覆盖的参考点 ID）；
- np.unique() 默认升序排序并去重，若需保持首次出现顺序，可改用 np.unique(..., return_index=True) 配合索引重排，但通常升序更符合分析直觉；
- 当 X 规模极大时，np.hstack 仍为内存友好操作（不创建中间 Python 列表），性能优于 [arr for arr in idxs if arr.size] + np.concatenate。
综上，np.unique(np.hstack(idxs)) 是处理 query_radius() 输出最简洁、高效、地道的 NumPy 方式，兼顾可读性与工程实践性。
