kd_tree.py

1"""2KD 树：建树、查询与维度灾难（第18章 §18.2.1）34正文用二维的七个点说明 KD 树的构造：5    A(50,50) B(10,70) C(80,15) D(25,20) E(40,85) F(70,85) G(10,60)6并给定划分规则：7    1) 在当前划分维度上取中位数作为划分点，保证左右子树尽量平衡；8    2) 划分维度按坐标轴轮流切换，任意连续 K 层内每个维度都用一次。9回溯剪枝的判据是10    |q_i − s| > r   →   兄弟子树不可能更近，剪掉1112本脚本：131) 先按正文的结构把七点 KD 树建出来（为保证与书中插图一致，根节点取 A 而不是中位数 E），14   打印每个节点的划分轴、划分值与超矩形范围；152) 实现通用 KD 树（中位数划分 + 叶子桶 + 超矩形下界），在合成数据上做精确 KNN，16   与暴力搜索逐条对比，确认它确实是"精确"搜索；173) 扫维度 d = 2 … 128，统计平均访问点数占全库的比例，复现维度灾难：18   d 小时只访问极少数点，d 大时几乎退化为全量扫描；194) 打印剪枝判据与超矩形距离公式的实现，便于对照正文。2021运行：python3 ch02_tree_based/kd_tree.py22"""2324from __future__ import annotations2526import pathlib27import sys2829import numpy as np3031sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))3233from common import knn_exact, make_queries, make_vectors, recall_at_k, table  # noqa: E4023435SEED_DIM = 7   # 维度扫描专用种子：保证每一维度都用同一份随机数序列3637NAMES = list("ABCDEFG")38PTS = np.array([39    [50, 50],   # A40    [10, 70],   # B41    [80, 15],   # C42    [25, 20],   # D43    [40, 85],   # E44    [70, 85],   # F45    [10, 60],   # G46], dtype=np.float32)4748# 正文插图的划分顺序：根 A(x)，左 B(y)，右 C(y)，再往下 G(x)49BOOK_SPLITS = {50    frozenset("ABCDEFG"): ("A", 0),51    frozenset("BDEG"): ("B", 1),52    frozenset("CF"): ("C", 1),53    frozenset("DG"): ("G", 0),54}555657class KDNode:58    """一个 KD 树节点：内部节点存划分轴 + 划分点，叶节点存样本；两者都存超矩形。"""5960    __slots__ = ("axis", "pivot", "left", "right", "pts", "lo", "hi", "depth")6162    def __init__(self) -> None:63        self.axis: int = -164        self.pivot: int = -165        self.left: KDNode | None = None66        self.right: KDNode | None = None67        self.pts: np.ndarray | None = None68        self.lo: np.ndarray | None = None69        self.hi: np.ndarray | None = None70        self.depth: int = 0717273def build_kd(X: np.ndarray, idx: np.ndarray, depth: int = 0,74             leaf_size: int = 8, book: bool = False,75             names: list[str] | None = None) -> KDNode:76    """77    递归建树。7879    book=True 时按正文插图的划分点（A → B/C → G）建树，否则严格用中位数。80    关键：叶节点保存点集，内部节点保存划分信息，两种节点都记录自己覆盖的超矩形，81    查询时用超矩形到查询点的距离做剪枝。82    """83    node = KDNode()84    node.depth = depth85    pts = X[idx]86    node.lo, node.hi = pts.min(0), pts.max(0)87    if book and names is not None:88        key = frozenset(names[i] for i in idx)89        if key in BOOK_SPLITS:90            name, axis = BOOK_SPLITS[key]91            pivot = names.index(name)92            side_l = [i for i in idx if i != pivot and X[i, axis] <= X[pivot, axis]]93            side_r = [i for i in idx if i != pivot and X[i, axis] > X[pivot, axis]]94            node.axis, node.pivot = axis, pivot95            node.left = build_kd(X, np.array(side_l, dtype=int), depth + 1, leaf_size, True, names) if side_l else None96            node.right = build_kd(X, np.array(side_r, dtype=int), depth + 1, leaf_size, True, names) if side_r else None97            return node98        node.pts = idx99        return node100101    if len(idx) <= leaf_size:102        node.pts = idx103        return node104    axis = depth % X.shape[1]105    order = idx[np.argsort(X[idx, axis], kind="stable")]106    mid = len(order) // 2107    node.axis, node.pivot = axis, order[mid]108    node.left = build_kd(X, order[:mid], depth + 1, leaf_size, book, names)109    node.right = build_kd(X, order[mid + 1:], depth + 1, leaf_size, book, names)110    return node111112113def dist_to_box(q: np.ndarray, lo: np.ndarray, hi: np.ndarray) -> float:114    """115    查询点到超矩形的最小距离（正文的逐维 δ 公式）：116        δ_i = l_i − q_i  (q_i < l_i) / 0 (在区间内) / q_i − u_i (q_i > u_i)117        d(q, R) = sqrt(Σ δ_i²)118    """119    delta = np.maximum(np.maximum(lo - q, q - hi), 0.0)120    return float(np.sqrt((delta * delta).sum()))121122123def print_book_tree(node: KDNode, X: np.ndarray, indent: str = "") -> None:124    """打印正文七点 KD 树的结构。"""125    if node.pts is not None:126        names = "".join(NAMES[i] for i in node.pts)127        print(f"{indent}叶节点 {{ {names} }}  超矩形 x∈[{node.lo[0]:.0f},{node.hi[0]:.0f}] "128              f"y∈[{node.lo[1]:.0f},{node.hi[1]:.0f}]")129        return130    axis, pivot = node.axis, node.pivot131    axis_name = "x" if axis == 0 else "y"132    print(f"{indent}内部节点 {NAMES[pivot]}({X[pivot, 0]:.0f},{X[pivot, 1]:.0f})  "133          f"按 {axis_name} = {X[pivot, axis]:.0f} 划分   "134          f"超矩形 x∈[{node.lo[0]:.0f},{node.hi[0]:.0f}] y∈[{node.lo[1]:.0f},{node.hi[1]:.0f}]")135    if node.left is not None:136        print_book_tree(node.left, X, indent + "  ├─L ")137    if node.right is not None:138        print_book_tree(node.right, X, indent + "  └─R ")139140141class KDTree:142    """带剪枝计数的通用 KD 树（中位数划分 + 叶子桶）。"""143144    def __init__(self, X: np.ndarray, leaf_size: int = 8) -> None:145        self.X = X146        self.root = build_kd(X, np.arange(len(X)), leaf_size=leaf_size)147        self.visited = 0          # 统计：本次查询访问到的点数148        self.nodes_visited = 0    # 统计：本次查询访问到的节点数149        self.total_visited = 0    # 统计：累计访问点数（跨多次查询）150        self.total_queries = 0    # 统计：累计查询次数151152    def _push(self, q: np.ndarray, i: int, heap: list, k: int) -> None:153        d = float(((self.X[i] - q) ** 2).sum())154        self.visited += 1155        if len(heap) < k:156            heap.append((d, i))157            heap.sort(reverse=True)158        elif d < heap[0][0]:159            heap[0] = (d, i)160            heap.sort(reverse=True)161162    def _search(self, node: KDNode, q: np.ndarray, heap: list, k: int) -> None:163        self.nodes_visited += 1164        if node.pts is not None:165            for i in node.pts:166                self._push(q, i, heap, k)167            return168169        self._push(q, node.pivot, heap, k)170        go_left = q[node.axis] <= self.X[node.pivot, node.axis]171        first, second = (node.left, node.right) if go_left else (node.right, node.left)172        if first is not None:173            self._search(first, q, heap, k)174        # 剪枝：当前最近邻距离 r，若兄弟超矩形距离 > r 则整棵子树都不可能更近175        if second is not None:176            ub = float("inf") if len(heap) < k else heap[0][0]177            if len(heap) < k or dist_to_box(q, second.lo, second.hi) ** 2 <= ub:178                self._search(second, q, heap, k)179180    def search(self, q: np.ndarray, k: int = 10) -> np.ndarray:181        heap: list = []182        self.visited = self.nodes_visited = 0183        self._search(self.root, q, heap, k)184        # 单次查询的计数在下一次调用时会被清零，跨查询的统计必须在这里累加185        self.total_visited += self.visited186        self.total_queries += 1187        return np.array([i for _, i in heap])188189190def main() -> None:191    print("KD 树：建树、查询与维度灾难（§18.2.1）\n")192193    # ---------------- 1) 正文的七点 KD 树 ----------------194    print("① 按正文插图重建七点 KD 树（根节点取 A，而不是排序后的中位数 E）：")195    names = NAMES196    root = build_kd(PTS, np.arange(7), book=True, names=names)197    print_book_tree(root, PTS)198    print("\n  正文的取舍：真正的中位数是 E(x=40)，但原论文插图以 A 为根。")199    print("  KD 树允许任意合法的轴划分，选谁只影响平衡性，不影响正确性；")200    print("  被选为划分点的样本放在内部节点上，不重复出现在子节点里。\n")201202    # ---------------- 2) 通用 KD 树 = 精确搜索 ----------------203    X, _ = make_vectors(n=2000, dim=8, n_clusters=16, spread=0.5)204    Q = make_queries(X, n_queries=100)205    k = 10206    tree = KDTree(X, leaf_size=8)207    truth_idx, _ = knn_exact(X, Q, k)208    pred = np.array([tree.search(q, k) for q in Q])209    avg_visited = tree.total_visited / len(Q)210    print("② 通用 KD 树（中位数划分、叶子桶 8）在 2000 条 8 维簇状数据上的表现：")211    print(f"  Recall@{k} = {recall_at_k(pred, truth_idx):.4f}  ← KD 树剪枝是"212          "「安全的」，只剪掉不可能更近的区域，因此是精确搜索")213    print(f"  平均访问点数 = {avg_visited:.0f} / {len(X)} "214          f"（{avg_visited / len(X):.1%} 的库）")215216    # 单条查询的剪枝过程217    q = Q[0]218    tree.search(q, k)219    print(f"\n  单条查询 q = [{', '.join(f'{v:.2f}' for v in q[:4])} …] 的过程：")220    print(f"    访问 {tree.nodes_visited} 个节点、{tree.visited} 个样本点，"221          f"得到 Top-{k}")222    print("    剪枝只发生在 dist_to_box(q, 兄弟超矩形) > 当前第 k 近距离时。\n")223224    # ---------------- 3) 维度灾难 ----------------225    # 关键：必须用"各向同性"（每个维度方差相同、无低维流形）的数据。226    # 簇状数据的中心只有 16 个，实际自由度是 16 + spread，维度升到 128 时227    # 数据仍躺在一个 16 维流形上，KD 树照样能剪枝——那测的是内在维度，不是维度灾难。228    print("③ 维度灾难：固定库规模 2000、叶子桶 8、Top-10，扫维度。")229    print("   左边是各向同性均匀数据（真正体现 d 的影响），右边是 16 簇数据（内在维度低）：")230    rows = []231    n_q = 60232    for dim in (2, 4, 8, 16, 32, 64, 128):233        rng = np.random.default_rng(SEED_DIM)234        Xu = rng.random((2000, dim), dtype=np.float32)                       # [0,1)^d235        Qu = (Xu[rng.integers(0, 2000, n_q)]236              + rng.normal(0.0, 0.05, (n_q, dim))).astype(np.float32)237        tu = KDTree(Xu, leaf_size=8)238        for qq in Qu:239            tu.search(qq, k)240        frac_u = tu.total_visited / (n_q * len(Xu))241242        Xc, _ = make_vectors(n=2000, dim=dim, n_clusters=16, spread=0.5,243                             seed=SEED_DIM)244        Qc = make_queries(Xc, n_queries=n_q, seed=SEED_DIM + 1)245        tc = KDTree(Xc, leaf_size=8)246        for qq in Qc:247            tc.search(qq, k)248        frac_c = tc.total_visited / (n_q * len(Xc))249250        rows.append([dim, f"{frac_u:.1%}", f"{frac_c:.1%}"])251    table(["维度 D", "均匀数据访问占比", "簇状数据访问占比"], rows)252    print("\n  读法：均匀数据在 D=8 时约 28%，D=16 已经 95%，D≥32 就是 100%——")253    print("  剪枝条件 dist(q,R) ≤ r 几乎对所有兄弟节点成立，KD 树退化成全量扫描。")254    print("  正文说的「D<20 才适用」在这组数字上非常准确。")255    print("  簇状数据始终停在 13% 附近，不是因为 KD 树变强了，而是因为数据实际")256    print("  只落在 16 个中心附近的低维流形上——起作用的是内在维度（intrinsic")257    print("  dimension），不是名义维度。真实 Embedding 往往介于两者之间：")258    print("  有簇结构，但簇内仍然很高维，所以工业上仍然不用 KD 树。\n")259260    # ---------------- 4) 回溯剪枝的具体判据 ----------------261    print("④ 两个必须记住的式子：")262    print("  划分：中位数 + 轴轮换 → 树高 O(log N)，建树 O(N log N)")263    print("  剪枝：|q_i − s| > r  ⟹  兄弟子树不可能更近（r 为当前第 k 近距离）")264    print("  更一般的形式是对整个超矩形求下界：dist(q,R) > r 才能剪。")265    print("  KD 树是精确搜索，它降的是平均访问量，不是精度；")266    print("  高维失效是因为 r 随 D 增大而变大，下界失去区分度，不是算法写错了。")267268269if __name__ == "__main__":270    main()