【堆】【kd-tree】bzoj2626 JZPFAR

#include<cstdio>
#include<queue>
#include<cstring>
#include<cmath>
#include<algorithm>
using namespace std;
typedef double db;
#define N 100001
#define EPS 0.0000001
#define INF 999999999999999999.0
#define KD 2//ά¶ÈÊý
int qp[KD];
int n,root,kd=2,K;
int dn;
struct Ans
{
    int p[KD],id;
    db d;
    Ans(){}
    Ans(int _p[],int _id,db _d){memcpy(p,_p,sizeof(p)); id=_id; d=_d;}
};
bool operator < (const Ans &a,const Ans &b)
{return fabs(a.d-b.d)>=EPS ? a.d>b.d : a.id<b.id;}
priority_queue<Ans>Heap;
db sqr(const int &x){return (db)x*(db)x;}
struct Node
{
    int minn[KD],maxx[KD],p[KD],id;
    int ch[2];
    void Init()
      {
        for(int i=0;i<kd;++i)
          minn[i]=maxx[i]=p[i];
      }
    db Dis()
      {
        db t=0;
        for(int i=0;i<kd;++i)
          {
            t+=sqr(max(0,qp[i]-minn[i]));
            t+=sqr(max(0,maxx[i]-qp[i]));
          }
        return sqrt(t);
      }
}T[N];
void Update(int rt)
{
    for(int i=0;i<2;++i)
      if(T[rt].ch[i])
        for(int j=0;j<kd;++j)
          {
            T[rt].minn[j]=min(T[rt].minn[j],T[T[rt].ch[i]].minn[j]);
            T[rt].maxx[j]=max(T[rt].maxx[j],T[T[rt].ch[i]].maxx[j]);
          }
}
db Dis(int a[],int b[])
{
    db t=0;
    for(int i=0;i<kd;++i)
      t+=sqr(a[i]-b[i]);
    return sqrt(t);
}
bool operator < (const Node &a,const Node &b){return a.p[dn]<b.p[dn];}
int Buildtree(int l=1,int r=n,int d=0)
{
    dn=d;
    int m=(l+r>>1);
    nth_element(T+l,T+m,T+r+1);
    T[m].Init();
    if(l!=m) T[m].ch[0]=Buildtree(l,m-1,(d+1)%kd);
    if(m!=r) T[m].ch[1]=Buildtree(m+1,r,(d+1)%kd);
    Update(m);
    return m;
}
void Query(int rt=root)
{
    db t=Dis(T[rt].p,qp);
    if(Heap.size()<K)
      Heap.push(Ans(T[rt].p,T[rt].id,t));
    else if(Heap.top().d-t<-EPS || (fabs(Heap.top().d-t)<EPS && T[rt].id<Heap.top().id))
      {
        Heap.pop();
        Heap.push(Ans(T[rt].p,T[rt].id,t));
      }
    db dd[2];
    for(int i=0;i<2;i++)
      if(T[rt].ch[i])
        dd[i]=T[T[rt].ch[i]].Dis();
      else dd[i]=-INF;
    bool f=(dd[0]>=dd[1]);
    if((dd[!f]-Heap.top().d>EPS || Heap.size()<K) && T[rt].ch[!f]) Query(T[rt].ch[!f]);
    if((dd[f]-Heap.top().d>EPS || Heap.size()<K) && T[rt].ch[f]) Query(T[rt].ch[f]);
}
int q;
int main()
{
//  freopen("bzoj2626.in","r",stdin);
//  freopen("bzoj3053.out","w",stdout);
    scanf("%d",&n);
    for(int i=1;i<=n;++i)
      {
        for(int j=0;j<kd;++j)
          scanf("%d",&T[i].p[j]);
        T[i].id=i;
      }
    Buildtree();
    root=(1+n>>1);
    scanf("%d",&q);
    for(;q;--q)
      {
        while(!Heap.empty())
          Heap.pop();
        for(int i=0;i<kd;++i)
          scanf("%d",&qp[i]);
        scanf("%d",&K);
        Query();
        printf("%d\n",Heap.top().id);
      }
    return 0;
}