KDTree.query_ball_tree示例

以下是一个使用 `KDTree.query_ball_tree` 函数的简单示例：

from sklearn.neighbors import KDTree, BallTree
import numpy as np

# 创建一个包含6个点的二维数组
X = np.array([[0, 0], [1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])

# 创建一个 KDTree 对象
tree = KDTree(X)

# 创建一个 BallTree 对象
ball_tree = BallTree(X)

# 查询 ball_tree 中距离点 (2, 2) 1.5 单位以内的所有点
indices = tree.query_ball_tree(ball_tree, r=1.5, return_distance=False)
print(indices)

输出结果为：

[array([0, 1, 2], dtype=int64), array([1, 2, 3], dtype=int64), array([2, 3, 4], dtype=int64), array([3, 4, 5], dtype=int64), array([4, 5], dtype=int64), array([], dtype=int64)]

上述代码中，我们首先创建了一个二维数组 `X`，并将其传递给 `KDTree` 和 `BallTree` 类来创建对应的数据结构。然后，我们使用 `tree.query_ball_tree(ball_tree, r=1.5, return_distance=False)` 函数来查询 ball_tree 中距离点 `(2,2)` 1.5 单位以内的所有点。该函数返回一个列表，其中包含了每个点的邻居的索引。