K维树 | 设置1(搜索和插入)

什么是K维树？

K维树又称KD树，是一种针对K维空间点的数据结构。它可以用于高维空间中的关键值存储和搜索，比如在搜索最近邻问题和桶归并问题上有着广泛的应用。K维树采用的是二叉搜索树的思想，其每个结点的值是一个K维向量，每个结点代表的是由一个超矩形（超立方体）划分的K维空间划分中的一个划分单元。

K维树的实现与操作

结点定义

在实现K维树之前，我们需要定义一个存储结点信息的结构体。结构体中包含了当前结点所属的点、结点深度、左右儿子等信息。具体代码实现如下：

struct kdtree{
    int p[N][K], lc[N], rc[N], sz[N], dep[N];
    int rt, tot;
    void pushup(int x){
        sz[x] = sz[lc[x]] + sz[rc[x]] + 1;
    }
    int build(int l, int r, int dep){
        if(l > r) return 0;
        int mid = l + r >> 1;
        int now = ++ tot;
        nth_element(p + l, p + mid, p + r + 1, cmp(dep));
        lc[now] = build(l, mid - 1, dep ^ 1);
        rc[now] = build(mid + 1, r, dep ^ 1);
        this -> dep[now] = dep;
        pushup(now);
        return now;
    }
}T;

插入操作

当我们需要在K维树中插入一个新结点时，我们需要利用K维树的性质完成。遍历每个结点的左右子树，根据所属深度的奇偶性判断要走到左子树还是右子树，并递归地继续查找。如果到达空结点，则新建一个结点并将其作为当前结点的左（右）儿子。

void insert(int nowp[]){
    int now = T.rt, dep = 0;
    while(1){
        ++ T.sz[now];
        if(cmp(dep)(nowp, T.p[now]) < 0){
            if(T.lc[now]) now = T.lc[now];
            else {
                T.lc[now] = ++ T.tot;
                for(int i = 0; i < K; ++ i) T.p[T.tot][i] = nowp[i];
                T.dep[T.tot] = dep + 1;
                ++ T.sz[T.tot];
                break;
            }
        }
        else {
            if(T.rc[now]) now = T.rc[now];
            else {
                T.rc[now] = ++ T.tot;
                for(int i = 0; i < K; ++ i) T.p[T.tot][i] = nowp[i];
                T.dep[T.tot] = dep + 1;
                ++ T.sz[T.tot];
                break;
            }
        }
        dep ^= 1;
    }
}

查找操作

查找操作需要遍历整个K维树，并根据当前结点所属深度与查找结点的大小关系判断要走到左子树还是右子树。一旦找到了目标结点，则将其返回。具体代码实现如下：

void ask(int now, const int nowp[], int &res){
    if(!now) return;
    res = min(res, dis(T.p[now], nowp));
    int dl = T.lc[now] ? dis(T.p[T.lc[now]], nowp) : INF;
    int dr = T.rc[now] ? dis(T.p[T.rc[now]], nowp) : INF;
    if(dl < dr){
        if(dl < res) ask(T.lc[now], nowp, res);
        if(dr < res) ask(T.rc[now], nowp, res);
    }
    else{
        if(dr < res) ask(T.rc[now], nowp, res);
        if(dl < res) ask(T.lc[now], nowp, res);
    }
}

参考代码

完整代码如下：