""" KD 树：建树、查询与维度灾难（第18章 §18.2.1） 正文用二维的七个点说明 KD 树的构造： A(50,50) B(10,70) C(80,15) D(25,20) E(40,85) F(70,85) G(10,60) 并给定划分规则： 1) 在当前划分维度上取中位数作为划分点，保证左右子树尽量平衡； 2) 划分维度按坐标轴轮流切换，任意连续 K 层内每个维度都用一次。 回溯剪枝的判据是 |q_i − s| > r → 兄弟子树不可能更近，剪掉 本脚本： 1) 先按正文的结构把七点 KD 树建出来（为保证与书中插图一致，根节点取 A 而不是中位数 E）， 打印每个节点的划分轴、划分值与超矩形范围； 2) 实现通用 KD 树（中位数划分 + 叶子桶 + 超矩形下界），在合成数据上做精确 KNN， 与暴力搜索逐条对比，确认它确实是"精确"搜索； 3) 扫维度 d = 2 … 128，统计平均访问点数占全库的比例，复现维度灾难： d 小时只访问极少数点，d 大时几乎退化为全量扫描； 4) 打印剪枝判据与超矩形距离公式的实现，便于对照正文。 运行：python3 ch02_tree_based/kd_tree.py """ from __future__ import annotations import pathlib import sys import numpy as np sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) from common import knn_exact, make_queries, make_vectors, recall_at_k, table # noqa: E402 SEED_DIM = 7 # 维度扫描专用种子：保证每一维度都用同一份随机数序列 NAMES = list("ABCDEFG") PTS = np.array([ [50, 50], # A [10, 70], # B [80, 15], # C [25, 20], # D [40, 85], # E [70, 85], # F [10, 60], # G ], dtype=np.float32) # 正文插图的划分顺序：根 A(x)，左 B(y)，右 C(y)，再往下 G(x) BOOK_SPLITS = { frozenset("ABCDEFG"): ("A", 0), frozenset("BDEG"): ("B", 1), frozenset("CF"): ("C", 1), frozenset("DG"): ("G", 0), } class KDNode: """一个 KD 树节点：内部节点存划分轴 + 划分点，叶节点存样本；两者都存超矩形。""" __slots__ = ("axis", "pivot", "left", "right", "pts", "lo", "hi", "depth") def __init__(self) -> None: self.axis: int = -1 self.pivot: int = -1 self.left: KDNode | None = None self.right: KDNode | None = None self.pts: np.ndarray | None = None self.lo: np.ndarray | None = None self.hi: np.ndarray | None = None self.depth: int = 0 def build_kd(X: np.ndarray, idx: np.ndarray, depth: int = 0, leaf_size: int = 8, book: bool = False, names: list[str] | None = None) -> KDNode: """ 递归建树。 book=True 时按正文插图的划分点（A → B/C → G）建树，否则严格用中位数。 关键：叶节点保存点集，内部节点保存划分信息，两种节点都记录自己覆盖的超矩形， 查询时用超矩形到查询点的距离做剪枝。 """ node = KDNode() node.depth = depth pts = X[idx] node.lo, node.hi = pts.min(0), pts.max(0) if book and names is not None: key = frozenset(names[i] for i in idx) if key in BOOK_SPLITS: name, axis = BOOK_SPLITS[key] pivot = names.index(name) side_l = [i for i in idx if i != pivot and X[i, axis] <= X[pivot, axis]] side_r = [i for i in idx if i != pivot and X[i, axis] > X[pivot, axis]] node.axis, node.pivot = axis, pivot node.left = build_kd(X, np.array(side_l, dtype=int), depth + 1, leaf_size, True, names) if side_l else None node.right = build_kd(X, np.array(side_r, dtype=int), depth + 1, leaf_size, True, names) if side_r else None return node node.pts = idx return node if len(idx) <= leaf_size: node.pts = idx return node axis = depth % X.shape[1] order = idx[np.argsort(X[idx, axis], kind="stable")] mid = len(order) // 2 node.axis, node.pivot = axis, order[mid] node.left = build_kd(X, order[:mid], depth + 1, leaf_size, book, names) node.right = build_kd(X, order[mid + 1:], depth + 1, leaf_size, book, names) return node def dist_to_box(q: np.ndarray, lo: np.ndarray, hi: np.ndarray) -> float: """ 查询点到超矩形的最小距离（正文的逐维 δ 公式）： δ_i = l_i − q_i (q_i < l_i) / 0 (在区间内) / q_i − u_i (q_i > u_i) d(q, R) = sqrt(Σ δ_i²) """ delta = np.maximum(np.maximum(lo - q, q - hi), 0.0) return float(np.sqrt((delta * delta).sum())) def print_book_tree(node: KDNode, X: np.ndarray, indent: str = "") -> None: """打印正文七点 KD 树的结构。""" if node.pts is not None: names = "".join(NAMES[i] for i in node.pts) print(f"{indent}叶节点 {{ {names} }} 超矩形 x∈[{node.lo[0]:.0f},{node.hi[0]:.0f}] " f"y∈[{node.lo[1]:.0f},{node.hi[1]:.0f}]") return axis, pivot = node.axis, node.pivot axis_name = "x" if axis == 0 else "y" print(f"{indent}内部节点 {NAMES[pivot]}({X[pivot, 0]:.0f},{X[pivot, 1]:.0f}) " f"按 {axis_name} = {X[pivot, axis]:.0f} 划分 " f"超矩形 x∈[{node.lo[0]:.0f},{node.hi[0]:.0f}] y∈[{node.lo[1]:.0f},{node.hi[1]:.0f}]") if node.left is not None: print_book_tree(node.left, X, indent + " ├─L ") if node.right is not None: print_book_tree(node.right, X, indent + " └─R ") class KDTree: """带剪枝计数的通用 KD 树（中位数划分 + 叶子桶）。""" def __init__(self, X: np.ndarray, leaf_size: int = 8) -> None: self.X = X self.root = build_kd(X, np.arange(len(X)), leaf_size=leaf_size) self.visited = 0 # 统计：本次查询访问到的点数 self.nodes_visited = 0 # 统计：本次查询访问到的节点数 self.total_visited = 0 # 统计：累计访问点数（跨多次查询） self.total_queries = 0 # 统计：累计查询次数 def _push(self, q: np.ndarray, i: int, heap: list, k: int) -> None: d = float(((self.X[i] - q) ** 2).sum()) self.visited += 1 if len(heap) < k: heap.append((d, i)) heap.sort(reverse=True) elif d < heap[0][0]: heap[0] = (d, i) heap.sort(reverse=True) def _search(self, node: KDNode, q: np.ndarray, heap: list, k: int) -> None: self.nodes_visited += 1 if node.pts is not None: for i in node.pts: self._push(q, i, heap, k) return self._push(q, node.pivot, heap, k) go_left = q[node.axis] <= self.X[node.pivot, node.axis] first, second = (node.left, node.right) if go_left else (node.right, node.left) if first is not None: self._search(first, q, heap, k) # 剪枝：当前最近邻距离 r，若兄弟超矩形距离 > r 则整棵子树都不可能更近 if second is not None: ub = float("inf") if len(heap) < k else heap[0][0] if len(heap) < k or dist_to_box(q, second.lo, second.hi) ** 2 <= ub: self._search(second, q, heap, k) def search(self, q: np.ndarray, k: int = 10) -> np.ndarray: heap: list = [] self.visited = self.nodes_visited = 0 self._search(self.root, q, heap, k) # 单次查询的计数在下一次调用时会被清零，跨查询的统计必须在这里累加 self.total_visited += self.visited self.total_queries += 1 return np.array([i for _, i in heap]) def main() -> None: print("KD 树：建树、查询与维度灾难（§18.2.1）\n") # ---------------- 1) 正文的七点 KD 树 ---------------- print("① 按正文插图重建七点 KD 树（根节点取 A，而不是排序后的中位数 E）：") names = NAMES root = build_kd(PTS, np.arange(7), book=True, names=names) print_book_tree(root, PTS) print("\n 正文的取舍：真正的中位数是 E(x=40)，但原论文插图以 A 为根。") print(" KD 树允许任意合法的轴划分，选谁只影响平衡性，不影响正确性；") print(" 被选为划分点的样本放在内部节点上，不重复出现在子节点里。\n") # ---------------- 2) 通用 KD 树 = 精确搜索 ---------------- X, _ = make_vectors(n=2000, dim=8, n_clusters=16, spread=0.5) Q = make_queries(X, n_queries=100) k = 10 tree = KDTree(X, leaf_size=8) truth_idx, _ = knn_exact(X, Q, k) pred = np.array([tree.search(q, k) for q in Q]) avg_visited = tree.total_visited / len(Q) print("② 通用 KD 树（中位数划分、叶子桶 8）在 2000 条 8 维簇状数据上的表现：") print(f" Recall@{k} = {recall_at_k(pred, truth_idx):.4f} ← KD 树剪枝是" "「安全的」，只剪掉不可能更近的区域，因此是精确搜索") print(f" 平均访问点数 = {avg_visited:.0f} / {len(X)} " f"（{avg_visited / len(X):.1%} 的库）") # 单条查询的剪枝过程 q = Q[0] tree.search(q, k) print(f"\n 单条查询 q = [{', '.join(f'{v:.2f}' for v in q[:4])} …] 的过程：") print(f" 访问 {tree.nodes_visited} 个节点、{tree.visited} 个样本点，" f"得到 Top-{k}") print(" 剪枝只发生在 dist_to_box(q, 兄弟超矩形) > 当前第 k 近距离时。\n") # ---------------- 3) 维度灾难 ---------------- # 关键：必须用"各向同性"（每个维度方差相同、无低维流形）的数据。 # 簇状数据的中心只有 16 个，实际自由度是 16 + spread，维度升到 128 时 # 数据仍躺在一个 16 维流形上，KD 树照样能剪枝——那测的是内在维度，不是维度灾难。 print("③ 维度灾难：固定库规模 2000、叶子桶 8、Top-10，扫维度。") print(" 左边是各向同性均匀数据（真正体现 d 的影响），右边是 16 簇数据（内在维度低）：") rows = [] n_q = 60 for dim in (2, 4, 8, 16, 32, 64, 128): rng = np.random.default_rng(SEED_DIM) Xu = rng.random((2000, dim), dtype=np.float32) # [0,1)^d Qu = (Xu[rng.integers(0, 2000, n_q)] + rng.normal(0.0, 0.05, (n_q, dim))).astype(np.float32) tu = KDTree(Xu, leaf_size=8) for qq in Qu: tu.search(qq, k) frac_u = tu.total_visited / (n_q * len(Xu)) Xc, _ = make_vectors(n=2000, dim=dim, n_clusters=16, spread=0.5, seed=SEED_DIM) Qc = make_queries(Xc, n_queries=n_q, seed=SEED_DIM + 1) tc = KDTree(Xc, leaf_size=8) for qq in Qc: tc.search(qq, k) frac_c = tc.total_visited / (n_q * len(Xc)) rows.append([dim, f"{frac_u:.1%}", f"{frac_c:.1%}"]) table(["维度 D", "均匀数据访问占比", "簇状数据访问占比"], rows) print("\n 读法：均匀数据在 D=8 时约 28%，D=16 已经 95%，D≥32 就是 100%——") print(" 剪枝条件 dist(q,R) ≤ r 几乎对所有兄弟节点成立，KD 树退化成全量扫描。") print(" 正文说的「D<20 才适用」在这组数字上非常准确。") print(" 簇状数据始终停在 13% 附近，不是因为 KD 树变强了，而是因为数据实际") print(" 只落在 16 个中心附近的低维流形上——起作用的是内在维度（intrinsic") print(" dimension），不是名义维度。真实 Embedding 往往介于两者之间：") print(" 有簇结构，但簇内仍然很高维，所以工业上仍然不用 KD 树。\n") # ---------------- 4) 回溯剪枝的具体判据 ---------------- print("④ 两个必须记住的式子：") print(" 划分：中位数 + 轴轮换 → 树高 O(log N)，建树 O(N log N)") print(" 剪枝：|q_i − s| > r ⟹ 兄弟子树不可能更近（r 为当前第 k 近距离）") print(" 更一般的形式是对整个超矩形求下界：dist(q,R) > r 才能剪。") print(" KD 树是精确搜索，它降的是平均访问量，不是精度；") print(" 高维失效是因为 r 随 D 增大而变大，下界失去区分度，不是算法写错了。") if __name__ == "__main__": main()