KD-Tree

Hardsoft

2025-12-11

学习记录

KD Tree

KD Tree 概述

目的：给定一些高维空间的点，构造一种数据结构，能够高效地对一个范围内的所有点进行查询和修改。
- 为了方便描述，我们定义当前空间是 $D$ 维的，每个点的信息形如 $ (p_1, p_2, \dots, p_D) $。
  询问的范围是 $ [(l_1, l_2, \dots, l_D), (r_1, r_2, \dots, r_D)] $,
  一个点 $p$ 在范围内，当且仅当 $ \forall i \in [1, D], l_i \le p_i \le r_i $。
  pushdown(u)：下传标记。
  pushup(u)：上传数据。
  fullin(x)：判断子树 $u$ 所有点是否在 $[L, R]$ 内。
  partin(x)：判断子树 $u$ 是否有点在 $[L, R]$ 内。
  selfin(x)：判断点 $u$ 是否在 $[L, R]$ 内。
  `
构造方法：
按照构造 平衡二叉搜索树 的规则，对于树中的每一层（深度相同的节点），以不同的维度切面上点的值（$p_d$）为关键字，建立平衡二叉搜索树。
- 一般的，我们使用以下方法设定每一层的“关键维度”：设当前深度为 $d$，以 $(d \bmod D + 1)$ 作为关键维度。即每次加深，维度切换 $1$。
- 实现中，使用 nth_element 求出当前位置应该放置的元素（即求关键字排序下的中位数）。
- 时间复杂度：$\mathcal O(n \log n)$，空间复杂度：$\mathcal O(n)$。

int build(int l = 1, int r = a[0], int k = 0) {
    int mid = l + r >> 1;
    std::nth_element(a+l, a+mid, a+r+1, [k](int i, int j){ return tr[i][k] < tr[j][k]; });
    int u = a[mid];
    tr[u].l = (l ^ mid) ? build(l, mid - 1, (k+1) % K) : 0;
    tr[u].r = (mid ^ r) ? build(mid + 1, r, (k+1) % K) : 0;
    pushup(u);
    return u;
}

查询/修改方法：
dfs 遍历每一层，
如果左/右区间一定不会修改或者一定不会有贡献，那么就不递归。
如果区间被包含，就打懒标记，结束递归。
类似于线段树。
- 时间复杂度：$\mathcal O(n^{2-\frac 1 D})$

LL que(int u) { // 查询 (L) 到 (R) 这个高维矩形内所有点的和
    if (!u) return 0;
    if (fullin(tr[u])) return tr[u].sum;
    if (!partin(tr[u])) return 0;
    LL ans = selfin(tr[u]) ? tr[u].w : 0;
    pushdown(u);
    return ans + que(tr[u].l) + que(tr[u].r);
}
void upd(int u) { // 根据 X 修改 (L) 到 (R) 这个高维矩形内所有点的信息（示例：+= X）
    if (!u) return;
    if (fullin(tr[u])) return tr[u].pull(X);
    if (!partin(tr[u])) return;
    if (selfin(tr[u])) tr[u].w += X;
    pushdown(u);
    upd(tr[u].l), upd(tr[u].r);
    pushup(u);
}

支持插入的 K-D Tree

替罪羊式维护 K-D Tree 是错误的解法！不保证时间复杂度。

我们考虑将整颗树拆成 $\mathcal O(\log n)$ 颗树构成的森林。
方便起见，我们设 $tr_i$ 表示大小 $\le 2^i$ 的树。
那么当我们插入一个元素 $p$ 的时候，我们枚举 $i=0,1,\dots k$：

如果 $size(tr_i) > 0$（即树不是空的），那么就把当前这颗树“拍扁”到数组 $a$，也就是恢复成散点的形式。

void flatten(int &u) { // 将 u 这颗子树拍扁，拼接到 a 数组的末尾
    if (!u) return;
    a[++a[0]] = u;
    pushdown(u);
    flatten(tr[u].l);
    flatten(tr[u].r);
    u = 0;
}

如果 $size(tr_i) = 0$（树是空的），那么就把存在 $a$ 数组当中的散点建树到 $tr_i$。

inline void insert(LL dat[], LL v) {
    ++tot;
    memcpy(tr[tot].dat, dat, sizeof tr[tot].dat);
    tr[tot].w = v;
    a[a[0] = 1] = tot;
    for (int i = 0; i < LOGN; i++) { // 如同二进制加法，每次加一
        if (!rt[i]) { rt[i] = build(); break; }
        else flatten(rt[i]);
    }
}

注意，如果要这样使用二进制分组的话，每次修改/查询也需要对每一颗子树进行操作，整体时间复杂度变为 $\mathcal O(n^{2 - \frac 1 D} \log n)$

inline LL query(LL l[], LL r[]) {
    memcpy(L.dat, l, sizeof L.dat);
    memcpy(R.dat, r, sizeof R.dat);
    LL ans = 0;
    for (int i = 0; i < LOGN; i++)
        ans += que(rt[i]);
    return ans;
}
inline void update(LL l[], LL r[], LL x) {
    X = x;
    memcpy(L.dat, l, sizeof L.dat);
    memcpy(R.dat, r, sizeof R.dat);
    for (int i = 0; i < LOGN; i++)
        upd(rt[i]);
}

有的时候，我们可以使用 K-D Tree 进行一些骗分，比如邻域查询（查询离点 $p$ 最近的点），但是这个时间复杂度是错的，慎用。

KD Tree 概述

支持 插入 的 K-D Tree

支持插入的 K-D Tree