UVALive - 7744 Finding Hotels - K-D Tree

#include <iostream>
#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
#define INF 1e18
const int maxn = 200000 + 7;
const int Size = 2;
struct Node{
   LL pos[Size];
   int id,price;
}node[maxn],op,point;
int split[maxn],now,n,m;
LL ans;
LL Distance(int a,int b){
    return (op.pos[b] - node[a].pos[b])*(op.pos[b] - node[a].pos[b]);
}
bool cmp(Node &a,Node &b){
   return a.pos[split[now]] < b.pos[split[now]];
}
int getSplit(int L,int R){
    double var = -1;
    int sp = 0;
     for(int i = 0;i<2;i++){
         double ave = 0;
         for(int j = L;j<=R;j++)ave+=node[j].pos[i];
         ave/=(R - L + 1);
         double sum = 0;
         for(int j = L;j<=R;j++)sum+=(node[j].pos[i] - ave)*(node[j].pos[i] - ave);
         sum/=(R - L + 1);
         if(sum > var){
            var = sum;
            sp = i;
         }
     }
     return sp;
}
void BuildTree(int L,int R){
     if(L > R)return;
     int mid = (L + R)>>1;
     split[now = mid] = getSplit(L,R);
     nth_element(node + L,node + mid,node + R + 1,cmp);
     BuildTree(L,mid-1);
     BuildTree(mid+1,R);
}
void Query(int L,int R){
    if(L > R)return;
    int mid = (L + R)>>1;
    LL dis = 0;
    for(int i = 0;i<2;i++)dis+=Distance(mid,i);
    if(dis < ans&&op.price >= node[mid].price){
        ans = dis;point = node[mid];
    }
    else if(dis==ans&&op.price >= node[mid].price&&node[mid].id < point.id){
        point = node[mid];
    }
    LL res = Distance(mid,split[mid]);
    if(op.pos[split[mid]] < node[mid].pos[split[mid]]){
        Query(L,mid-1);
        if(res<ans)Query(mid+1,R);
    }
    else{
        Query(mid+1,R);
        if(res<ans)Query(L,mid-1);
    }
}
int main()
{
    int T;
    scanf("%d",&T);
    while(T--){
        scanf("%d%d",&n,&m);
        for(int i = 1;i<=n;i++){
            scanf("%lld%lld%d",&node[i].pos[0],&node[i].pos[1],&node[i].price);
            node[i].id = i;
        }
        BuildTree(1,n);
        int k = m;
        while(k--){
            scanf("%lld%lld%d",&op.pos[0],&op.pos[1],&op.price);
            ans = INF;
            Query(1,n);
            printf("%lld %lld %d\n",point.pos[0],point.pos[1],point.price);
        }
    }
    return 0;
}