K-D Tree 模版
K-D Tree 模版
下面这个模版（洛谷P1429为例）可以实现k维空间查询给定点的m邻近点，存储于que这个大根堆中。
#include<bits/stdc++.h>
#define ll long long
#define sq(x) (x)*(x)
#define N (250000)
using namespace std;
int idx, k, n, m, q;//idx为某个维度的编号 k为维度个数
struct Node
{
double x[5];
bool operator < (const Node &u) const
{
return x[idx] < u.x[idx];//以全局变量idx对应的维度为排序标准
}
} P[N];
typedef pair<double, Node> PDN;//用于优先队列的排序
priority_queue<PDN> que;
struct KD_Tree
{
int sz[N << 2];
Node p[N << 2];
void build (int i, int l, int r, int dep) {//dep初始为0 l初始为0 r初始为n - 1
if (l > r) return;
int mid = (l + r) >> 1;
idx = dep % k;
sz[i] = r - l;
sz[i << 1] = sz[i << 1 | 1] = -1;
nth_element(P + l, P + mid, P + r + 1);//O(n)选取中位数
p[i] = P[mid];
build(i << 1, l, mid - 1, dep + 1);
build(i << 1 | 1, mid + 1, r, dep + 1);
}
void query (int i, int m, int dep, Node a) {//查询点为a dep为深度 m为最近邻m个
if (sz[i] == -1) return;
PDN tmp = PDN(0, p[i]);
for(int j = 0; j < k; j++)
tmp.first += 1.0 * sq(1.0 * tmp.second.x[j] - 1.0 * a.x[j]);//这里是对应的欧式距离 如果要用闵可夫斯基距离的话需要手动修改
tmp.first = sqrt(tmp.first);
int lc = i << 1,rc = i << 1 | 1, dim = dep % k, flag = 0;
if (a.x[dim] >= p[i].x[dim]) swap(lc, rc);
//如果sz[lc]为-1，~sz[lc]为0，就不进行搜索了
if (~sz[lc]) query(lc, m, dep + 1, a);
if (que.size() < m) que.push(tmp), flag = 1;
else {
if (tmp.first < que.top().first) que.pop(), que.push(tmp);
//它的兄弟所在的区域里面可能会有比它距离查询点更近的点，所以要对它的兄弟也进行搜索
if (sqrt(1.0 * sq(1.0 * a.x[dim] - 1.0 * p[i].x[dim])) < que.top().first) flag = 1;
}
if (~sz[rc] && flag) query(rc, m, dep + 1, a);
}
} KDT;
int main() {
k = 2;
cin >> n;
for(int i = 0; i < n; i++) {
for(int j = 0; j < k; j++) {
cin >> P[i].x[j];
}
}
KDT.build(1, 0, n - 1, 0);
double ans = 1e18;
m = 2;//最近的点 但会把自己算上（自己到自己的距离为0）因此m取2
for(int i = 0; i < n; i++) {
KDT.query(1, m, 0, P[i]);
if(que.top().first < ans) {
ans = que.top().first;
}
while(que.size()) que.pop();
}
printf("%.4lf", ans);
return 0;
}
还有一道比较好的题是洛谷P4357，查询平面第K远点对。一种做法是首先找到全局的最远点对(x, y)，然后从这个点对的一个点y出发找到除(x, y)的最远点对(y, z)，如此进行K次（有点类似树的直径的两次BFS），去重使用unordered_map。但上面的模版无法实现查询最远点对。题解的思路是启发式搜索降低复杂度，对于kdt的每个节点维护其超矩形边界（实际上就是二维，因为算法竞赛中基本都是二维）。等有时间可以尝试将上述模版改成可以查询平面k近/k远点对～下面先放上zero4438大佬的模版：
#include<iostream>
#include<cstdio>
#include<algorithm>
#include<unordered_map>
#include<queue>
#define ll long long
using namespace std;
const int maxn=1e5+10;
int n,k;
int comp;
struct point
{
int num;
ll d[2];
friend bool operator <(point a,point b){return a.d[comp]<b.d[comp];}
friend ll dis(point a,point b){return (a.d[0]-b.d[0])*(a.d[0]-b.d[0])+(a.d[1]-b.d[1])*(a.d[1]-b.d[1]);}//统一把根号去掉
};
unordered_map<int,bool>ban[maxn];//记录每个元素已经和谁配对
struct par
{
point x,y;ll dis;
friend bool operator <(par a,par b){return a.dis<b.dis;}
};
priority_queue<par>q;
struct kdtree
{
struct node
{
point p;ll minv[2],maxv[2];int l,r;
};
node t[maxn];
point sta[maxn];
int rt,cnt;
par res;
void update(int u)
{
int l=t[u].l,r=t[u].r;
for(int i=0;i<=1;i++)
{
t[u].minv[i]=t[u].maxv[i]=t[u].p.d[i];
if(l)t[u].minv[i]=min(t[u].minv[i],t[l].minv[i]),t[u].maxv[i]=max(t[u].maxv[i],t[l].maxv[i]);
if(r)t[u].minv[i]=min(t[u].minv[i],t[r].minv[i]),t[u].maxv[i]=max(t[u].maxv[i],t[r].maxv[i]);
}
}
int build(int l,int r,int f)
{
int mid=(l+r)>>1;
comp=f;
nth_element(sta+l,sta+mid,sta+r+1);
int u=++cnt;
t[u].p=sta[mid];
if(l<mid)t[u].l=build(l,mid-1,f^1);
if(r>mid)t[u].r=build(mid+1,r,f^1);
update(u);return u;
}
ll getmax(int u,point x)
{
ll ret=0;
for(int i=0;i<=1;i++)ret+=max(x.d[i]-t[u].minv[i],t[u].maxv[i]-x.d[i])*max(x.d[i]-t[u].minv[i],t[u].maxv[i]-x.d[i]);
return ret;
}
void query(int u,point x)//常规 K-D tree 计算最远点
{
if(ban[x.num].find(t[u].p.num)==ban[x.num].end()&&dis(t[u].p,x)>res.dis)//只有当该点没有与询问点配过对时才更新答案
{
res.y=t[u].p;res.dis=dis(t[u].p,x);
}
int l=t[u].l,r=t[u].r;
ll dl=0,dr=0;
if(l)dl=getmax(l,x);if(r)dr=getmax(r,x);
if(dl>dr)
{
if(dl>res.dis)query(l,x);if(dr>res.dis)query(r,x);
}
else
{
if(dr>res.dis)query(r,x);if(dl>res.dis)query(l,x);
}
}
}tr;
par ans;
int main()
{
scanf("%d%d",&n,&k);k*=2;//令 k=k*2
for(int i=1;i<=n;i++)scanf("%lld%lld",&tr.sta[i].d[0],&tr.sta[i].d[1]),tr.sta[i].num=i;
tr.rt=tr.build(1,n,0);
for(int i=1;i<=n;i++)
{
tr.res.x=tr.res.y=tr.sta[i];tr.res.dis=0;
tr.query(tr.rt,tr.sta[i]);//计算最远点
q.push(tr.res); //放入大根堆
ban[tr.res.x.num][tr.res.y.num]=1;//记录已出现点对
}
while(k--)
{
par now=q.top();q.pop();//取出当前最远点
ans=now;
tr.res.x=tr.res.y=now.x;tr.res.dis=0;
tr.query(tr.rt,now.x);//放入次远点
q.push(tr.res);ban[tr.res.x.num][tr.res.y.num]=1;
}
printf("%lld",ans.dis);
return 0;
}
