[KD TREE] BZOJ 3053 The Closest M Points

k维的kdtree啊

怎么说呢 kdtree就是各种剪枝啊

#include<cstdio>
#include<cstdlib>
#include<algorithm>
#include<queue>
#define inf 2147483647
using namespace std;
 
inline char nc(){
  static char buf[100000],*p1=buf,*p2=buf;
  if (p1==p2) { p2=(p1=buf)+fread(buf,1,100000,stdin); if (p1==p2) return EOF; }
  return *p1++;
}
 
inline void read(int &x){
  char c=nc(),b=1;
  for (;!(c>='0' && c<='9');c=nc()) if (c=='-') b=-1;
  for (x=0;c>='0' && c<='9';x=x*10+c-'0',c=nc()); x*=b;
}
 
inline int sqr(int x){ return x*x; }
 
const int N=50005,M=5;
 
int n,m,D;
 
struct node{
  int l,r;
  int d[M],mx[M],mn[M];
  int &operator [](int x){ return d[x]; }
  friend bool operator == (node A,node B){ int f=1; for (int i=0;i<m;i++) f&=(A[i]==B[i]); return f; }
  friend bool operator <  (node A,node B){ return A[D]<B[D]; }
  friend int dist(node A,node B) { int ret=0; for (int i=0;i<m;i++) ret+=sqr(A[i]-B[i]); return ret; }
}tmp[N];
 
struct abcd{
  node a; int d;
  abcd(){}
  abcd(node _a,int _d){ a=_a; d=_d; }
  bool operator < (const abcd B) const{ return d<B.d; }
}ans[15];
int K;
 
struct KDT{
  node T[N],now;
  int root;
  void update(int x){
    int l=T[x].l,r=T[x].r;
    for (int i=0;i<m;i++)
    {
      T[x].mn[i]=T[x].mx[i]=T[x][i];
      if(l) T[x].mn[i]=min(T[x].mn[i],T[l].mn[i]),T[x].mx[i]=max(T[x].mx[i],T[l].mx[i]);
      if(r) T[x].mn[i]=min(T[x].mn[i],T[r].mn[i]),T[x].mx[i]=max(T[x].mx[i],T[r].mx[i]);
    }
  }
  int Reb(int l,int r,int D){
    if (l>r) return 0;
    int mid=(l+r)>>1;
    ::D=D; nth_element(tmp+l,tmp+mid,tmp+r+1);
    T[mid]=tmp[mid];
    T[mid].l=Reb(l,mid-1,(++D)%m);
    T[mid].r=Reb(mid+1,r,(++D)%m);
    update(mid); return mid;
  }
  int Get(node A){
    int ret=0;
    for (int i=0;i<m;i++)
      if (!(A.mn[i]<=now[i] && now[i]<=A.mx[i]))
        ret+=min(sqr(A.mn[i]-now[i]),sqr(A.mx[i]-now[i]));
    return ret;
  }
  void Query(int x){
    int tmp=dist(T[x],now);
    ans[K+1]=abcd(T[x],tmp); int pnt=K+1;
    while (pnt>1 && ans[pnt]<ans[pnt-1]) swap(ans[pnt],ans[pnt-1]),pnt--;
    int l=T[x].l,r=T[x].r,dl=inf,dr=inf;
    if (l) dl=Get(T[l]);
    if (r) dr=Get(T[r]);
    if (dl<dr){
      if (dl<ans[K].d) Query(l);
      if (dr<ans[K].d) Query(r);
    }else{
      if (dr<ans[K].d) Query(r);
      if (dl<ans[K].d) Query(l);
    }
  }
}KD;
 
int main()
{
  int Q;
  while (~scanf("%d%d",&n,&m))
  {
    // read(n); read(m);
    for (int i=1;i<=n;i++)
      for (int j=0;j<m;j++)
      {
        scanf("%d",&tmp[i][j]);
        // read(tmp[i][j]);
        tmp[i].mn[j]=tmp[i].mx[j]=tmp[i][j];
      }
    KD.root=KD.Reb(1,n,0);
    // read(Q);
    scanf("%d",&Q);
    while (Q--)
    {
      for (int i=0;i<m;i++)
        //read(KD.now[i]);
        scanf("%d",&KD.now[i]);
      // read(K);
      scanf("%d",&K);
      for (int i=1;i<=K;i++) ans[i].d=inf;
      KD.Query(KD.root);
      printf("the closest %d points are:\n",K);
      for (int i=1;i<=K;i++)
        for (int j=0;j<m;j++)
          printf("%d%c",ans[i].a[j],j==m-1?'\n':' ');
    }
  }
  return 0;
}