求完整的Python版BSP树（Binary Space Partitioning Tree）实现示例

嘿，我太懂你找合适BSP树实现的痛苦了——这类能直接用于最近邻（NN）搜索的完整实现确实不像KD树那么遍地都是。刚好我整理了一个适配NN搜索场景的Python实现，而且是轴对齐的版本，和你现有的KD树结构接近，方便你调整成一致的配置做对比：

import numpy as np
from scipy.spatial.distance import euclidean

class BSPNode:
    def __init__(self):
        self.left = None  # 左子树：分割面左侧的点集
        self.right = None  # 右子树：分割面右侧的点集
        self.split_axis = None  # 分割轴（比如0对应x轴，1对应y轴）
        self.split_value = None  # 分割轴上的分割阈值
        self.points = None  # 叶子节点存储的点集（非叶子节点为None）

def build_bsp_tree(points, depth=0, split_strategy='cycle'):
    """
    构建BSP树，支持两种分割轴选择策略
    :param points: numpy数组，形状为(样本数, 特征数)
    :param depth: 当前树的深度，用于循环选轴策略
    :param split_strategy: 分割轴选择方式：'cycle'循环选轴 / 'random'随机选轴
    :return: BSP树的根节点
    """
    n_samples, n_features = points.shape
    
    # 当点数小于等于阈值时，作为叶子节点存储点集
    if n_samples <= 4:
        leaf_node = BSPNode()
        leaf_node.points = points
        return leaf_node
    
    # 选择分割轴
    if split_strategy == 'cycle':
        split_axis = depth % n_features  # 和KD树一致的循环选轴
    elif split_strategy == 'random':
        split_axis = np.random.randint(0, n_features)
    else:
        raise ValueError("分割策略仅支持'cycle'或'random'")
    
    # 按分割轴排序，取中位数作为分割值（保证树的平衡）
    sorted_indices = np.argsort(points[:, split_axis])
    median_idx = n_samples // 2
    split_value = points[sorted_indices[median_idx], split_axis]
    
    # 分割点集
    left_points = points[points[:, split_axis] <= split_value]
    right_points = points[points[:, split_axis] > split_value]
    
    # 递归构建左右子树
    current_node = BSPNode()
    current_node.split_axis = split_axis
    current_node.split_value = split_value
    current_node.left = build_bsp_tree(left_points, depth + 1, split_strategy)
    current_node.right = build_bsp_tree(right_points, depth + 1, split_strategy)
    
    return current_node

def bsp_nearest_neighbor(node, query_point, best_point=None, best_dist=np.inf):
    """
    BSP树的最近邻搜索，带剪枝逻辑
    :param node: 当前遍历的BSP节点
    :param query_point: 查询点，numpy数组
    :param best_point: 当前找到的最近点
    :param best_dist: 当前找到的最近距离
    :return: 最近点、最近距离
    """
    # 叶子节点：遍历所有点找最近邻
    if node.points is not None:
        for point in node.points:
            current_dist = euclidean(query_point, point)
            if current_dist < best_dist:
                best_dist = current_dist
                best_point = point
        return best_point, best_dist
    
    # 计算查询点到分割面的距离（轴对齐场景下就是分割轴上的差值绝对值）
    query_axis_val = query_point[node.split_axis]
    split_val = node.split_value
    
    # 优先遍历查询点所在的子树
    if query_axis_val <= split_val:
        best_point, best_dist = bsp_nearest_neighbor(node.left, query_point, best_point, best_dist)
        # 如果查询点到分割面的距离小于当前最佳距离，说明另一侧可能有更近的点，需要遍历
        if abs(query_axis_val - split_val) < best_dist:
            best_point, best_dist = bsp_nearest_neighbor(node.right, query_point, best_point, best_dist)
    else:
        best_point, best_dist = bsp_nearest_neighbor(node.right, query_point, best_point, best_dist)
        if abs(query_axis_val - split_val) < best_dist:
            best_point, best_dist = bsp_nearest_neighbor(node.left, query_point, best_point, best_dist)
    
    return best_point, best_dist

# 测试用例
if __name__ == "__main__":
    # 生成2维随机测试点
    np.random.seed(42)
    test_points = np.random.rand(100, 2)
    query = np.array([0.5, 0.5])
    
    # 构建循环选轴的BSP树（和KD树选轴逻辑一致）
    bsp_root = build_bsp_tree(test_points, split_strategy='cycle')
    
    # 搜索最近邻
    nearest_point, nearest_dist = bsp_nearest_neighbor(bsp_root, query)
    print(f"查询点: {query}")
    print(f"最近邻点: {nearest_point}")
    print(f"最近距离: {nearest_dist:.4f}")

关键说明和适配建议： #

• 这个实现是轴对齐BSP树，和KD树的结构高度相似，你只需要调整几个参数就能和你的KD树配置对齐：
  • 叶子节点的点数阈值（当前是4），改成和你的KD树一致；
  • 分割策略默认是cycle（循环选轴），和标准KD树逻辑完全匹配；
• 如果你需要测试更通用的BSP树（非轴对齐），可以修改分割面的表示方式（比如用超平面系数a1*x1 + a2*x2 + ... + d = 0），但这样和KD树的对比会更复杂，建议先从轴对齐版本开始；
• 关于你之前修改KD树失败的问题：轴对齐BSP树和KD树的核心差异仅在于分割轴的选择灵活性——KD树严格循环选轴，而BSP树可以选任意轴（甚至非轴对齐平面）。如果你的KD树是自平衡的，也可以给这个实现加上平衡逻辑（比如左右子树高度差过大时重新分割）。

备注：内容来源于stack exchange，提问作者Lucas Duan