本文介绍如何将 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() 输出的标准范式：简洁、健壮、零循环、完全向量化，适用于任何规模的邻域索引聚合任务。