KDtree

[•ELHoria]

Salut ma poate ajuta cineva cu problema http://infoarena.ro/problema/votare.
Am implementat kd treeul asa

Cod:

#include <cstdio>
#include <utility>
#define MP make_pair
#define PII pair<int,int>
#define x first
#define y second

using namespace std;

const long long INF = 1ll<<61;

long long squared_distance(const PII &A,const PII &B)
{
	return 1ll*(A.x - B.x)*(A.x - B.x) + 1ll*(A.y - B.y)*(A.y - B.y);
}

class KDnode{
	public :
		KDnode *left_son , *right_son , *parent;
		bool axis;
		PII point;
		KDnode(bool axisc,KDnode *node,const PII &p) {
			axis = axisc;
			left_son = right_son = NULL;
			parent = node;
			point = p;
		}
		void print();
		bool is_leaf() {return (left_son == NULL && right_son == NULL);}
		bool insert(const PII &p);
		bool erase(const PII &p);
		KDnode* find_replacement(bool w_axis,const int &tip);
		KDnode* find_node(const PII &p);
		long long nearest_distance(const PII &p,long long best);
};

long long KDnode::nearest_distance(const PII &p,long long best)
{
	if(this->axis == 0 ? p.x < this->point.x : p.y < this->point.y) 
	{
		if(this->left_son != NULL && squared_distance(p,left_son->point) <= best)
			best = left_son->nearest_distance(p, p == left_son->point ? best : squared_distance(p,left_son->point));
	
		if(this->right_son != NULL && squared_distance(p,right_son->point) <= best)
			best = right_son->nearest_distance(p,p == right_son->point ? best : squared_distance(p,right_son->point));
	}
	else
	{
		if(this->right_son != NULL && squared_distance(p,right_son->point) <= best)
			best = right_son->nearest_distance(p,p == right_son->point ? best : squared_distance(p,right_son->point));
	
		if(this->left_son != NULL && squared_distance(p,left_son->point) <= best)
			best = left_son->nearest_distance(p, p == left_son->point ? best : squared_distance(p,left_son->point));
	}
	return best;
}

void KDnode::print() 
{
	if(this == NULL) return;
	left_son->print();
	printf("%d %d\n",this->point.x,this->point.y);
	right_son->print();
}

KDnode* KDnode::find_node(const PII &p)
{
	if(this->point == p) return this;
	if((this->axis == 0 ? p.x < this->point.x : p.y < this->point.y)) {
			if(left_son != NULL) return left_son->find_node(p);
		}
	else
	{
		if(right_son != NULL) 
			return right_son->find_node(p);
	}
	return NULL;
}

bool KDnode::insert(const PII &p)
{
	if(this->point == p) return false;

	if((this->axis == 0 ? p.x < this->point.x : p.y < this->point.y) ) {
			if(left_son == NULL) { 
				left_son = new KDnode(!this->axis,this,p);
				return true;
			}
			return left_son->insert(p);
		}
	else
	{
		if(right_son == NULL) { 
			right_son = new KDnode(!this->axis,this,p);
			return true;
		}
		return right_son->insert(p);
	}
	return false;
}

KDnode* KDnode::find_replacement(bool w_axis,const int &tip)
{
	KDnode *R , *best = NULL;
	if(this->is_leaf()) return this;
	if(tip != -1) best = this;
	if(this->right_son != NULL)
	{
		R = right_son->find_replacement(w_axis,1);
		if(tip == -1) best = R;
		else
		if(w_axis == 0) 
		{
			if(R->point.x > best->point.x) best = R;
		}
		else
		if(w_axis == 1) 
		{
			if(R->point.y > best->point.y) best = R;
		}
	}
	else
	if(this->left_son != NULL) {
		R = left_son->find_replacement(w_axis,0);
		if(tip == -1) best = R;
		else
		if(w_axis == 0) 
		{
			if(R->point.x < best->point.x) best = R;
		}
		else
		if(w_axis == 1)
			if(R->point.y < best->point.y) best = R;
	}
	return best;
}
bool KDnode::erase(const PII &p)
{
	if(this == NULL) return false;
	if(is_leaf()) 
	{
		KDnode *father = this->parent;
		if(parent != NULL) {
		if(father->left_son != NULL && father->left_son->point == p) father->left_son = NULL;
		else 
		if(father->right_son != NULL && father->right_son->point == p) father->right_son = NULL;
		delete this;
		}
	}
	else
	{
		KDnode *R = this->find_replacement(this->axis,-1);
		this->point = R->point;
		R->erase(R->point);
	}
	return true;
}

class KDtree{ 
	public:
	KDnode *root;
	KDtree() {
	root = NULL;
	}
	bool add(const PII &p);
	bool remove(const PII &p);
	void print_tree();
	long long nearest_neighbour(const PII &p);
};

void KDtree::print_tree()
{
	if(root == NULL) { 
		printf("EMPTY TREE\n");
		return;
	}
	root->print();
	printf("\n");
}

bool KDtree::add(const PII &p)
{
	if(root == NULL) {
		root = new KDnode(0,NULL,p);
		return true;
	}
	KDnode *tar = root->find_node(p);
	if(tar == NULL)	return root->insert(p);
	return 1;
}

bool KDtree::remove(const PII &p)
{
	if(root == NULL) return 0;
	KDnode *tar = root->find_node(p);
	if(tar != NULL) tar->erase(p);
	return true;
}

long long KDtree::nearest_neighbour(const PII &p)
{
	if(root == NULL) return -1;
	if(p != root->point) return root->nearest_distance(p,squared_distance(p,root->point));
	return root->nearest_distance(p,INF);
}

int main()
{
	freopen("votare.in","r",stdin);
	freopen("votare.out","w",stdout);
	int N , t;
	PII c_point;
	KDtree *tree = new KDtree;
	for(scanf("%d",&N);N;N--)
	{
		scanf("%d %d %d",&t,&c_point.x,&c_point.y);
		if(t == 0) tree->add(c_point);
		else
		if(t == 1) tree->remove(c_point); 
		else
		if(t == 2) printf("%lld\n",tree->nearest_neighbour(c_point));
	}
	return 0;
}

Merge bine pentru exemple mici , stergerea si inserarea de noduri merg bine  , banuiesc ca problema e la query pentru nearest neighbor , daca ma poate ajuta cineva sa imi zica daca fac gresit la nearest neighbor.