本文介绍如何将 scikit-learn KDTree.query_radius() 返回的嵌套数组（含空数组）一键展平并去重，避免显式循环，直接获得全局唯一索引数组。
本文介绍如何将 scikit-learn `kdtree.query_radius()` 返回的嵌套数组（含空数组）一键展平并去重，避免显式循环，直接获得全局唯一索引数组。
在使用 sklearn.neighbors.KDTree 进行邻域搜索时，query_radius() 方法会为查询点集 Y 中的每个点返回一个一维 NumPy 数组（dtype=int64），表示在半径 r 内的所有参考点索引。由于不同查询点的邻域大小不一，输出是一个 object 类型的 NumPy 数组，其每个元素是长度可变的 ndarray——其中部分可能为空（如 array([], dtype=int64)），这给后续统一处理（如去重、统计、索引提取）带来不便。
直接对 idxs 调用 np.concatenate(idxs) 会报错（因 concatenate 不支持 object 数组中含空数组），而手动遍历 + extend 则违背向量化初衷。更简洁、高效且符合 NumPy 惯例的方案是使用 np.hstack()：
from sklearn.neighbors import KDTree import numpy as np # 示例数据（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) # 输出：object array of ndarrays # ✅ 一行完成：展平 + 去重 unique_indices = np.unique(np.hstack(idxs)) print(unique_indices) # 输出类似：[ 0 1 2 ... 47]（长度 ≤ 49 的升序唯一整数）
⚠️ 注意事项：
- np.hstack() 要求输入为序列（如 list 或 tuple），但 idxs 是 object 数组，恰好可被 hstack 直接接受（它会自动迭代各元素并水平拼接）；
- 空数组 [] 在 hstack 中被安全忽略，无需预过滤；
- np.unique() 默认升序排序并去重，返回一维 int64 数组；若需保留原始顺序，可改用 np.unique(..., return_index=True) 配合索引重排（但通常升序更符合分析需求）；
- 若 X 和 Y 维度较高或数据量极大，该方法仍保持内存友好性——hstack 是视图友好的拼接，unique 也基于高效哈希实现。
综上，np.unique(np.hstack(idxs)) 是处理 query_radius() 输出的标准范式：简洁、健壮、零循环、完全向量化，适用于任何规模的邻域索引聚合任务。
