WireRouting/include/KDtree.h

#include "stdafx.h"
#include "BVH.h"
#include<cmath>
#include<vector>
#include<queue>
#include<map>
#include<iostream>
#include<fstream>
#include<sstream>
using namespace std;


struct point{
	double x[3];
	int id;
	point(){
		x[0]=x[1]=x[2]=0;
		id=0;
	}
	//friend bool operator < (const point &a,const point &b)
	//	{return a.x[cmpid]<b.x[cmpid];}
};
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];
}


struct node{
	double dis;
	int id;
	bool operator < (node b)const {
		return dis<b.dis||(dis==b.dis&&id<b.id);
	}
	node(){
		dis=0;
		id=0;
	}
	node(int id1,double dis1){
		dis=dis1;
		id=id1;
	}
};
struct tree{
	point p;double mx[3],mn[3];
	int ls,rs,id;
	tree(){
		ls=rs=id=0;
	}
};
struct KDtree{
	int n,tot,root;
	double X,Y,Z;
	
	priority_queue<node> q;
	vector<int> vec;

	
	point p[MAXPointNum];
	tree t[MAXPointNum];

	KDtree(){
		X=Y=Z=0;
		n=tot=root=0;
	}
	inline double 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;
	}
	inline double 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);
	}
	inline void 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 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 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);
		}
	}
	inline void 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 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;
	}
	inline void build(){
		cout<<root<<"  "<<n<<endl;
		build(root,1,n,0);
	}

	//������s������k�����ı���
	inline vector<int> 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���ڵ����е��ı���
	inline vector<int> search_by_dis(P& s,double d){
		X=s.x;Y=s.y;Z=s.z;
		vec.clear();
		queryd(root,d);
		return vec;
	}
}kdtree;