#include "pch.h" // use stdafx.h in Visual Studio 2017 and earlier
#include "KDtree.h"


bool cmp0(const point& a, const point& b) {
	return a.x[0] < b.x[0];
}
bool cmp1(const point& a, const point& b) {
	return a.x[1] < b.x[1];
}
bool cmp2(const point& a, const point& b) {
	return a.x[2] < b.x[2];
}


point::point(){
	x[0]=x[1]=x[2]=0;
	id=0;
}


bool node::operator < (node b)const {
	return dis<b.dis||(dis==b.dis&&id<b.id);
}
node::node(){
	dis=0;
	id=0;
}
node::node(int id1,double dis1){
	dis=dis1;
	id=id1;
}

tree::tree(){
	ls=rs=id=0;
}

KDtree::KDtree(){
	X=Y=Z=0;
	n=tot=root=0;
}


double KDtree::dis(tree& x) {
	double P = (x.p.x[0] - X) * (x.p.x[0] - X);
	double Q = (x.p.x[1] - Y) * (x.p.x[1] - Y);
	double O = (x.p.x[2] - Z) * (x.p.x[2] - Z);
	return P + Q + O;
}
double KDtree::mndis(tree& x) {
	double P = (x.mn[0] - X) * (x.mn[0] - X);
	double M = (x.mx[0] - X) * (x.mx[0] - X);
	if (X >= x.mn[0] && X <= x.mx[0])P = M = 0;
	double Q = (x.mn[1] - Y) * (x.mn[1] - Y);
	double N = (x.mx[1] - Y) * (x.mx[1] - Y);
	if (Y >= x.mn[1] && Y <= x.mx[1])Q = N = 0;
	double O = (x.mn[2] - Z) * (x.mn[2] - Z);
	double U = (x.mx[2] - Z) * (x.mx[2] - Z);
	if (Z >= x.mn[2] && Z <= x.mx[2])O = U = 0;
	return min(P, M) + min(Q, N) + min(O, U);
}
void KDtree::update(int x) {
	if (!x)  return;
	int l = t[x].ls, r = t[x].rs;
	if (l) t[x].mn[0] = min(t[x].mn[0], t[l].mn[0]),
		t[x].mn[1] = min(t[x].mn[1], t[l].mn[1]),
		t[x].mn[2] = min(t[x].mn[2], t[l].mn[2]),
		t[x].mx[0] = max(t[x].mx[0], t[l].mx[0]),
		t[x].mx[1] = max(t[x].mx[1], t[l].mx[1]),
		t[x].mx[2] = max(t[x].mx[2], t[l].mx[2]);
	if (r) t[x].mn[0] = min(t[x].mn[0], t[r].mn[0]),
		t[x].mn[1] = min(t[x].mn[1], t[r].mn[1]),
		t[x].mn[2] = min(t[x].mn[2], t[r].mn[2]),
		t[x].mx[0] = max(t[x].mx[0], t[r].mx[0]),
		t[x].mx[1] = max(t[x].mx[1], t[r].mx[1]),
		t[x].mx[2] = max(t[x].mx[2], t[r].mx[2]);
}

void KDtree::query(int x){
	if(!x) return ;
	double res=dis(t[x]);
	if(res<q.top().dis||(res==q.top().dis&&t[x].id<q.top().id)) 
		q.pop(),q.push(node(t[x].id,res));
	int l=t[x].ls,r=t[x].rs;
	double ld=0,rd=0;
	if(l) ld=mndis(t[l]);
	if(r) rd=mndis(t[r]);
	//cout<<x<<"  "<<ld<<"  "<<rd<<"    "<<res<<"  "<<q.top().id<<" "<<q.top().dis<<endl;
	if(ld<rd){
		if(ld<=q.top().dis) query(l);
		if(rd<=q.top().dis) query(r);
	}
	else{
		if(rd<=q.top().dis) query(r);
		if(ld<=q.top().dis) query(l);
	}
}

void KDtree::queryd(int x,double d){
	if(!x) return ;
	double res=dis(t[x]);
	//cout<<t[x].p.id<<"  "<<res<<endl;
	if(res<=d*d)vec.push_back(t[x].p.id);
	int l=t[x].ls,r=t[x].rs;
	double ld=0,rd=0;
	if(l) ld=mndis(t[l]);
	if(r) rd=mndis(t[r]);
	//cout<<x<<"  "<<ld<<"  "<<rd<<"    "<<res<<"  "<<q.top().id<<" "<<q.top().dis<<endl;
	if(ld<rd){
		if(ld<=d*d) queryd(l,d);
		if(rd<=d*d) queryd(r,d);
	}
	else{
		if(rd<=d*d) queryd(r,d);
		if(ld<=d*d) queryd(l,d);
	}
}	

void KDtree::add(P& s) {
	n++;
	p[n].x[0] = s.x;
	p[n].x[1] = s.y;
	p[n].x[2] = s.z;
	p[n].id = n;

}

void KDtree::build(int &x,int l,int r,int k){
	if(l>r) return ;
	x=++tot;
	int mid=(l+r)>>1;
	//cout<<"stage1  "<<x<<endl;
	if (k == 0)nth_element(p + l, p + mid, p + r + 1, cmp0);
	else if (k == 1)nth_element(p + l, p + mid, p + r + 1, cmp1);
	else if (k == 2)nth_element(p + l, p + mid, p + r + 1, cmp2);
	//cout<<"stage2"<<endl;
	t[x].p=p[mid];t[x].id=t[x].p.id;
	t[x].mn[0]=t[x].mx[0]=t[x].p.x[0];
	t[x].mn[1]=t[x].mx[1]=t[x].p.x[1];
	t[x].mn[2]=t[x].mx[2]=t[x].p.x[2];
	build(t[x].ls,l,mid-1,(k+1)%3);
	build(t[x].rs,mid+1,r,(k+1)%3);
	//cout<<"stage3"<<endl;
	update(x);
	//cout<<x<<"  "<<t[x].ls<<"  "<<t[x].rs<<"           "<<t[x].mx[0]<<" "<<t[x].mx[1]<<" "<<t[x].mx[2]<<endl;
}


void KDtree::build() {
	cout << root << "  " << n << endl;
	build(root, 1, n, 0);
}

vector<int> KDtree::search_by_k(P& s, int k) {
	X = s.x; Y = s.y; Z = s.z;
	while (q.size()) q.pop();
	for (int j = 1; j <= k; j++) q.push(node(0, 1e9));
	query(root);
	vector<int> veck;
	while (q.size()) {
		if (q.top().id != 0)veck.push_back(q.top().id);
		q.pop();
	}
	return veck;
}

//返回距离s点d以内的所有点的编号
vector<int> KDtree::search_by_dis(P& s, double d) {
	X = s.x; Y = s.y; Z = s.z;
	vec.clear();
	queryd(root, d);
	return vec;
}