[BZOJ3053][KD树]The Closest M Points

题意

---

给定k维空间的n个点，每次询问给出一个点，求于给出的点距离最小的m个点

---

KD树标准模板题

#include <cstdio>
#include <iostream>
#include <algorithm>
#include <queue>
#define N 100010

using namespace std;

typedef pair<int,int> pairs;

inline int abs(int a){return a<0?-a:a;}
inline int sqr(int x){return x*x;}

int id,K;
struct node{
  int a[6],mn[6],mx[6],g,l,r;
  int &operator [](int x){
    return a[x];
  }
  friend bool operator <(node a,node b){
    return a[id]<b[id];
  }
  inline void print(){
    printf("%d",a[1]);
    for(int i=2;i<=K;i++) printf(" %d",a[i]);
    putchar('\n');
  }
}Point[N];

priority_queue<pairs> Q;

struct kd_tree{
  node p[N];
  int rt,q[N],cnt;
  inline void updat(int x){
    for(int i=1;i<=K;i++){
      p[x].mn[i]=p[x].mx[i]=p[x][i];
      if(p[x].l)
    p[x].mn[i]=min(p[x].mn[i],p[p[x].l].mn[i]),
      p[x].mx[i]=max(p[x].mx[i],p[p[x].l].mx[i]);
      if(p[x].r)
    p[x].mn[i]=min(p[x].mn[i],p[p[x].r].mn[i]),
      p[x].mx[i]=max(p[x].mx[i],p[p[x].r].mx[i]);

    }
  }
  inline int build(int l,int r,int t){
    if(l>r) return 0; id=t;
    int mid=l+r>>1;
    nth_element(Point+l,Point+mid,Point+r+1);
    p[mid]=Point[mid];
    p[mid].l=build(l,mid-1,t%K+1);
    p[mid].r=build(mid+1,r,t%K+1);
    updat(mid); return mid;
  }
  inline int dist(int g,node x){
    int r=0;
    if(g==0) return 1<<30;
    for(int i=1;i<=K;i++) r+=sqr(p[g][i]-x[i]);
    return r;
  }
  inline int wnt(int g,node x){
    int r=0;
    for(int i=1;i<=K;i++){
      if(p[g].mn[i]>x[i]) r+=sqr(p[g].mn[i]-x[i]);
      if(p[g].mx[i]<x[i]) r+=sqr(p[g].mx[i]-x[i]);
    }
    return r;
  }
  void query(int g,node x){
    if(!g) return ;
    int d=dist(g,x),dl=wnt(p[g].l,x),dr=wnt(p[g].r,x);
    if(d<Q.top().first) Q.pop(),Q.push(pairs(d,g));
    if(dl<dr){
      if(dl<Q.top().first) query(p[g].l,x);
      if(dr<Q.top().first) query(p[g].r,x);
    }
    else{
      if(dr<Q.top().first) query(p[g].r,x);
      if(dl<Q.top().first) query(p[g].l,x);
    }
  }
  inline void query(node x,int w){
    for(int i=1;i<=w;i++) Q.push(pairs(1<<30,0));
    query(rt,x); cnt=0;
    while(!Q.empty()){
      q[++cnt]=Q.top().second;
      Q.pop();
    }
    printf("the closest %d points are:\n",w);
    for(int i=cnt;i;i--) p[q[i]].print();
  }
}KDT;

int n,m;

int main(){
  while(~scanf("%d%d",&n,&K)){
    for(int i=1;i<=n;i++)
      for(int j=1;j<=K;j++) scanf("%d",&Point[i][j]);
    KDT.rt=KDT.build(1,n,1);
    scanf("%d",&m);
    for(int i=1;i<=m;i++){
      node x; int y;
      for(int i=1;i<=K;i++) scanf("%d",&x[i]);
      scanf("%d",&y);
      KDT.query(x,y);
    }
  }
}