KDT 学习笔记

· · 个人记录

介绍

K-D Tree，一种可以高效维护 k 维空间信息的数据结构，貌似可以理解为 k 维的平衡树。

主要思想

后面我们将认为 k 为常数，且 k \ge 2，在 OI 的多数题目中，有 k = 2。

K-D Tree 只支持建树后查询，不支持插入，但是可以通过二进制分组的方式实现 \mathcal{O}(n \log ^ 2 n) 的插入。

树的形态

K-D Tree 为二叉树，其每一个节点都对应了一个元素，每一棵子树都对应了一个 k 维矩阵，且树高为严格的 \log n + O(1)。

建树

递归建树，假设当前需要将 a_{l \sim r} 建树。

首先我们需要确定一个维度，找到这个维度在 a_{l \sim r} 中的中位数，作为当前节点，把在该维度上比中位数小的放左边，大的放右边，递归下去建树即可。

使用 nth_element 寻找中位数可以使时间复杂度正确。

为保证时间复杂度正确，我们需要轮流划分这 k 个维度。

查询

由于每个节点都对应了一个 k 维矩阵，所以像线段树一样查询即可。

时间复杂度

详见 OI-wiki，建树复杂度为 \mathcal{O}(n \log n)，查询复杂度为 \mathcal{O}(n ^ {1 - \frac{1}{k}})。

二进制分组

二进制分组可以使 K-D Tree 支持插入，具体地，建 \mathcal{O}(\log n) 棵大小为 2 的幂次的 K-D Tree，每次插入时建一棵大小为 1 的 K-D Tree，将大小相同的合并即可。

实现上，可找到第一个大小不存在的树，然后将比它小的树都拍到序列上，再建树。

查询时可在每棵树上都进行查询，时间复杂度不变。

势能分析，二进制分组后插入的时间复杂度为 \mathcal{O}(n \log ^ 2 n)。

例题

模板题，将节点插入再查询区间和即可。

代码

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

const int N = 5e5 + 5;
const int B = 21;

int n, cnt, tot, root[B], a[N], val[N], sum[N], lson[N], rson[N], L[N][2], R[N][2], x[N][2];

int ql[2], qr[2], ans;

inline void pushup(int u) {
    sum[u] = sum[lson[u]] + sum[rson[u]] + val[u];
    for (int k = 0; k < 2; ++k) {
        L[u][k] = R[u][k] = x[u][k];
        if (lson[u]) {
            L[u][k] = min(L[u][k], L[lson[u]][k]);
            R[u][k] = max(R[u][k], R[lson[u]][k]);
        }
        if (rson[u]) {
            L[u][k] = min(L[u][k], L[rson[u]][k]);
            R[u][k] = max(R[u][k], R[rson[u]][k]);
        }
    }
}

inline int build(int l, int r, int dep) {
    int mid = (l + r) >> 1;
    nth_element(a + l, a + mid, a + r + 1, [&](int u, int v) {
        return x[u][dep] < x[v][dep];
    });
    int p = a[mid];
    if (l < mid) {
        lson[p] = build(l, mid - 1, dep ^ 1);
    } else {
        lson[p] = 0;
    }
    if (mid < r) {
        rson[p] = build(mid + 1, r, dep ^ 1);
    } else {
        rson[p] = 0;
    }
    pushup(p);
    return p;
}

inline int query(int u) {
    bool flg = true;
    for (int k = 0; k < 2; ++k) {
        flg &= (ql[k] <= L[u][k] && R[u][k] <= qr[k]);
    }
    if (flg) {
        return sum[u];
    }
    flg = false;
    for (int k = 0; k < 2; ++k) {
        flg |= (ql[k] > R[u][k] || qr[k] < L[u][k]);
    }
    if (flg) {
        return 0;
    }
    int res = 0;
    flg = true;
    for (int k = 0; k < 2; ++k) {
        flg &= (ql[k] <= x[u][k] && x[u][k] <= qr[k]);
    }
    if (flg) {
        res = val[u];
    }
    res += query(lson[u]) + query(rson[u]);
    return res;
}

inline void append(int u) {
    if (!u) {
        return;
    }
    a[++cnt] = u;
    append(lson[u]);
    append(rson[u]);
}

int main() {
    #ifdef LOCAL
        assert(freopen("test.in", "r", stdin));
        assert(freopen("test.out", "w", stdout));
    #endif
    ios::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
    cin >> n;
    while (true) {
        int op;
        cin >> op;
        if (op == 1) {
            ++tot;
            cin >> x[tot][0] >> x[tot][1] >> val[tot];
            x[tot][0] ^= ans, x[tot][1] ^= ans, val[tot] ^= ans;
            a[cnt = 1] = tot;
            for (int i = 0; i < B; ++i) {
                if (!root[i]) {
                    root[i] = build(1, cnt, 0);
                    break;
                } else {
                    append(root[i]);
                    root[i] = 0;
                }
            }
        } else if (op == 2) {
            cin >> ql[0] >> ql[1] >> qr[0] >> qr[1];
            ql[0] ^= ans, ql[1] ^= ans, qr[0] ^= ans, qr[1] ^= ans;
            ans = 0;
            for (int i = 0; i < B; ++i) {
                if (root[i]) {
                    ans += query(root[i]);
                }
            }
            cout << ans << "\n";
        } else {
            break;
        }
    }
    return 0;
}

其它用途

因为 K-D Tree 每个子树对应了一个 k 维矩形，所以可以求类似平面最近点对之类的问题。

枚举平面中每一个点，在 K-D Tree 上 DFS 查询，若查询节点到当前子树对应的矩形的最近距离比答案大，则剪枝。

同理，可以将查询节点到当前子树对应的矩形的最近距离作为估价函数，采用启发式递归左右儿子会更优。

但最坏时间复杂度仍是 \mathcal{O}(n ^ 2)，随机数据下可以很快。