#ifndef _KD_TREE_H_
#define _KD_TREE_H_
#include <memory>
#include <vector>
#include <algorithm>
#include <iostream>
#include <functional>
#include <iomanip>
#include <stack>
#include <array>
#include <cfloat>
#include <cmath>
namespace zstd
{
struct threeD_node
{
double value[];//value[0] = x, value[1] = y, value[2] = z
threeD_node()
{
value[] = 0.0;
value[] = 0.0;
value[] = 0.0;
}
threeD_node(double x, double y, double z)
{
value[] = x;
value[] = y;
value[] = z;
}
};
struct sort_for_threeD_node
{
int dimension;
sort_for_threeD_node(int d) :dimension(d){}
bool operator ()(const threeD_node& lhs, const threeD_node& rhs)
{
if (dimension == )
return lhs.value[] < rhs.value[];
else if (dimension == )
return lhs.value[] < rhs.value[];
else if (dimension == )
return lhs.value[] < rhs.value[];
else
std::cerr << "error in sort_for_threeD_node"<< std::endl;
return false;
}
};
struct kd_node
{
double value[];//value[0] = x, value[1] = y, value[2] = z
int kv;//0=x, 1=y, 2=z
bool is_leaf;
kd_node *left, *right;
kd_node()
{
value[] = value[] = value[] = 0.0;
kv = -;
is_leaf = false;
left = nullptr;
right = nullptr;
}
kd_node(const kd_node& node)
{
value[] = node.value[];
value[] = node.value[];
value[] = node.value[];
kv = node.kv;
is_leaf = node.is_leaf;
left = node.left;
right = node.right;
}
kd_node& operator = (const kd_node& node)
{
value[] = node.value[];
value[] = node.value[];
value[] = node.value[];
kv = node.kv;
is_leaf = node.is_leaf;
left = node.left;
right = node.right;
return *this;
}
};
class kd_tree
{
private:
std::shared_ptr<kd_node> root;
std::vector<threeD_node>& vec_ref;
const int k = ;
const int cspace = ;
private:
int get_dimension(int n) const
{
return n % k;
}
void sort_by_dimension(std::vector<threeD_node>& v, int dimension, int l, int r);
kd_node* build_tree(int left, int right, kd_node* sp_node, int dimension);
void _print_tree(kd_node* sp, bool left, int space);
double distance(const kd_node& lhs, const threeD_node& rhs);
public:
explicit kd_tree(std::vector<threeD_node>&);
kd_tree(const kd_tree&) = delete;
kd_tree operator = (const kd_tree&) = delete;
~kd_tree(){};
void print_tree();
std::vector<threeD_node> find_k_nearest(int k, const threeD_node& D);
};
void kd_tree::sort_by_dimension(std::vector<threeD_node>& v, int dimension, int l, int r)
{
sort_for_threeD_node s(dimension);
std::sort(v.begin()+l, v.begin()+r, s);
}
kd_tree::kd_tree(std::vector<threeD_node>& v) :vec_ref(v)
{
if (vec_ref.empty())
root = nullptr;
else
{
root = std::make_shared<kd_node>();
int dimension = ;
sort_by_dimension(vec_ref, dimension, , vec_ref.size());
int mid = vec_ref.size() / ;
root->value[] = vec_ref[mid].value[];
root->value[] = vec_ref[mid].value[];
root->value[] = vec_ref[mid].value[];
root->kv = dimension;
if (vec_ref.size() == )//root is leaf
{
root->left = nullptr;
root->right = nullptr;
root->is_leaf = true;
}
else
{
root->is_leaf = false;
root->left = build_tree(, mid - , root->left, get_dimension(dimension + ));
root->right = build_tree(mid + , vec_ref.size() - , root->right, get_dimension(dimension + ));
}
}
}
kd_node* kd_tree::build_tree(int left, int right, kd_node* sp_node, int dimension)
{
dimension = get_dimension(dimension);
sort_by_dimension(vec_ref, dimension, left, right + );
if(left == right)//leaf
{
sp_node = new kd_node();
sp_node->value[] = vec_ref[left].value[];
sp_node->value[] = vec_ref[left].value[];
sp_node->value[] = vec_ref[left].value[];
sp_node->kv = dimension;
sp_node->is_leaf = true;
sp_node->left = nullptr;
sp_node->right = nullptr;
return sp_node;
}
else if (left < right)
{
int mid = left + (right - left) / ;
sp_node = new kd_node();
sp_node->value[] = vec_ref[mid].value[];
sp_node->value[] = vec_ref[mid].value[];
sp_node->value[] = vec_ref[mid].value[];
sp_node->kv = dimension;
sp_node->is_leaf = false;
sp_node->left = nullptr;
sp_node->right = nullptr;
sp_node->left = build_tree(left, mid - , sp_node->left, get_dimension(dimension + ));
sp_node->right = build_tree(mid + , right, sp_node->right, get_dimension(dimension + ));
return sp_node;
}
return nullptr;
}
void kd_tree::_print_tree(kd_node* sp, bool left, int space)
{
if (sp != nullptr)
{
_print_tree(sp->right, false, space + cspace);
std::cout << std::setw(space);
std::cout << "(" <<
sp->value[] << ", " <<
sp->value[] << ", " <<
sp->value[] << ")";
if (left)
std::cout << "left";
else
std::cout << "right";
if (sp->is_leaf)
std::cout << "------leaf";
std::cout << std::endl;
_print_tree(sp->left, true, space + cspace);
}
else
std::cout << std::endl;
}
void kd_tree::print_tree()
{
std::cout << "kd_tree : " << std::endl;
if (root != nullptr)
{
int space = ;
_print_tree(root->right, false, space + cspace);
std::cout << "(" <<
root->value[] << ", " <<
root->value[] << ", " <<
root->value[] << ")root" << std::endl;
_print_tree(root->left, true, space + cspace);
}
}
double kd_tree::distance(const kd_node& lhs, const threeD_node& rhs)
{
double v0 = lhs.value[] - rhs.value[];
double v1 = lhs.value[] - rhs.value[];
double v2 = lhs.value[] - rhs.value[];
return sqrt(v0 * v0 + v1 * v1 + v2 * v2);
}
std::vector<threeD_node> kd_tree::find_k_nearest(int ks, const threeD_node& D)
{
std::vector<threeD_node> res;
const kd_node *ptr_kd_node;
if (static_cast<std::size_t>(ks) > vec_ref.size())
return res;
std::stack<kd_node> s;
struct pair
{
double distance;
kd_node node;
pair() :distance(DBL_MAX), node(){ }
bool operator < (const pair& rhs)
{
return distance < rhs.distance;
}
};
std::unique_ptr<pair[]> ptr_pair(new pair[ks]);
//pair *ptr_pair = new pair[ks]();
if (!ptr_pair)
exit(-);
if (!root)//the tree is empty
return std::vector<threeD_node>();
else
{
if (D.value[root->kv] < root->value[root->kv])
{
s.push(*root);
ptr_kd_node = root->left;
}
else
{
s.push(*root);
ptr_kd_node = root->right;
}
while (ptr_kd_node != nullptr)
{
if (D.value[ptr_kd_node->kv] < ptr_kd_node->value[ptr_kd_node->kv])
{
s.push(*ptr_kd_node);
ptr_kd_node = ptr_kd_node->left;
}
else
{
s.push(*ptr_kd_node);
ptr_kd_node = ptr_kd_node->right;
}
}
while (!s.empty())
{
kd_node popped_kd_node;//±£´æ×îеĴÓÕ»ÖÐpop³öµÄkd_node
popped_kd_node = s.top();
s.pop();
double dist = distance(popped_kd_node, D);
std::sort(&ptr_pair[], &ptr_pair[ks]);
if (dist < ptr_pair[ks-].distance)
{
ptr_pair[ks-].distance = dist;
ptr_pair[ks-].node = popped_kd_node;
}
if (abs(D.value[popped_kd_node.kv] - popped_kd_node.value[popped_kd_node.kv])
>= dist)//Ô²²»ºÍpopped_kd_nodeµÄÁíÒ»°ëÇøÓòÏཻ
continue;
else//Ô²ºÍpopped_kd_nodeµÄÁíÒ»°ëÇøÓòÏཻ
{
if (D.value[popped_kd_node.kv] < popped_kd_node.value[popped_kd_node.kv])//right
{
kd_node *ptr = popped_kd_node.right;
while (ptr != nullptr)
{
s.push(*ptr);
if (D.value[ptr->kv] < ptr->value[ptr->kv])
ptr = ptr->left;
else
ptr = ptr->right;
}
}
else//left
{
kd_node *ptr = popped_kd_node.left;
while (ptr != nullptr)
{
s.push(*ptr);
if (D.value[ptr->kv] < ptr->value[ptr->kv])
ptr = ptr->left;
else
ptr = ptr->right;
}
}
}
}//end of while
for(int i = ; i != ks; ++i)
res.push_back(threeD_node(ptr_pair[i].node.value[],
ptr_pair[i].node.value[], ptr_pair[i].node.value[]));
}//end of else
//delete ptr_pair;
return res;
}
}//end of namespace zstd
#endif
