added kdtree implementation
This commit is contained in:
parent
a2ba8f14d3
commit
98cc39ffd8
70
data_structures/kdtree/build.hpp
Normal file
70
data_structures/kdtree/build.hpp
Normal file
@ -0,0 +1,70 @@
|
|||||||
|
#ifndef MEMGRAPH_DATA_STRUCTURES_KDTREE_BUILD_HPP
|
||||||
|
#define MEMGRAPH_DATA_STRUCTURES_KDTREE_BUILD_HPP
|
||||||
|
|
||||||
|
#include <vector>
|
||||||
|
#include <algorithm>
|
||||||
|
#include <functional>
|
||||||
|
|
||||||
|
#include "math.hpp"
|
||||||
|
#include "kdnode.hpp"
|
||||||
|
|
||||||
|
namespace kd {
|
||||||
|
|
||||||
|
template <class T, class U>
|
||||||
|
using Nodes = std::vector<KdNode<T, U>*>;
|
||||||
|
|
||||||
|
template <class T, class U>
|
||||||
|
KdNode<T, U>* build(Nodes<T, U>& nodes, byte axis = 0)
|
||||||
|
{
|
||||||
|
// if there are no elements left, we've completed building of this branch
|
||||||
|
if(nodes.empty())
|
||||||
|
return nullptr;
|
||||||
|
|
||||||
|
// comparison function to use for sorting the elements
|
||||||
|
auto fsort = [axis](KdNode<T, U>* a, KdNode<T, U>* b) -> bool
|
||||||
|
{ return kd::math::axial_distance(a->coord, b->coord, axis) < 0; };
|
||||||
|
|
||||||
|
size_t median = nodes.size() / 2;
|
||||||
|
|
||||||
|
// partial sort nodes vector to compute median and ensure that elements
|
||||||
|
// less than median are positioned before the median so we can slice it
|
||||||
|
// nicely
|
||||||
|
|
||||||
|
// internal implementation is O(n) worst case
|
||||||
|
// tl;dr http://en.wikipedia.org/wiki/Introselect
|
||||||
|
std::nth_element(nodes.begin(), nodes.begin() + median, nodes.end(), fsort);
|
||||||
|
|
||||||
|
// set axis for the node
|
||||||
|
auto node = nodes.at(median);
|
||||||
|
node->axis = axis;
|
||||||
|
|
||||||
|
// slice the vector into two halves
|
||||||
|
auto left = Nodes<T, U>(nodes.begin(), nodes.begin() + median);
|
||||||
|
auto right = Nodes<T, U>(nodes.begin() + median + 1, nodes.end());
|
||||||
|
|
||||||
|
// recursively build left and right branches
|
||||||
|
node->left = build(left, axis ^ 1);
|
||||||
|
node->right = build(right, axis ^ 1);
|
||||||
|
|
||||||
|
return node;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <class T, class U, class It>
|
||||||
|
KdNode<T, U>* build(It first, It last)
|
||||||
|
{
|
||||||
|
Nodes<T, U> kdnodes;
|
||||||
|
|
||||||
|
std::transform(first, last, std::back_inserter(kdnodes),
|
||||||
|
[&](const std::pair<Point<T>, U>& element) {
|
||||||
|
auto key = element.first;
|
||||||
|
auto data = element.second;
|
||||||
|
return new KdNode<T, U>(key, data);
|
||||||
|
});
|
||||||
|
|
||||||
|
// build the tree from the kdnodes and return the root node
|
||||||
|
return build(kdnodes);
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif
|
48
data_structures/kdtree/kdnode.hpp
Normal file
48
data_structures/kdtree/kdnode.hpp
Normal file
@ -0,0 +1,48 @@
|
|||||||
|
#ifndef GEOAPI_KDTREE_KDNODE_HPP
|
||||||
|
#define GEOAPI_KDTREE_KDNODE_HPP
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
|
||||||
|
#include "point.hpp"
|
||||||
|
|
||||||
|
namespace kd {
|
||||||
|
|
||||||
|
template <class T, class U>
|
||||||
|
class KdNode
|
||||||
|
{
|
||||||
|
public:
|
||||||
|
KdNode(const U& data)
|
||||||
|
: axis(0), coord(Point<T>(0, 0)), left(nullptr), right(nullptr), data(data) { }
|
||||||
|
|
||||||
|
KdNode(const Point<T>& coord, const U& data)
|
||||||
|
: axis(0), coord(coord), left(nullptr), right(nullptr), data(data) { }
|
||||||
|
|
||||||
|
KdNode(unsigned char axis, const Point<T>& coord, const U& data)
|
||||||
|
: axis(axis), coord(coord), left(nullptr), right(nullptr), data(data) { }
|
||||||
|
|
||||||
|
KdNode(unsigned char axis, const Point<T>& coord, KdNode<T, U>* left, KdNode<T, U>* right, const U& data)
|
||||||
|
: axis(axis), coord(coord), left(left), right(right), data(data) { }
|
||||||
|
|
||||||
|
~KdNode();
|
||||||
|
|
||||||
|
unsigned char axis;
|
||||||
|
|
||||||
|
Point<T> coord;
|
||||||
|
|
||||||
|
KdNode<T, U>* left;
|
||||||
|
KdNode<T, U>* right;
|
||||||
|
|
||||||
|
U data;
|
||||||
|
};
|
||||||
|
|
||||||
|
template <class T, class U>
|
||||||
|
KdNode<T, U>::~KdNode()
|
||||||
|
{
|
||||||
|
delete left;
|
||||||
|
delete right;
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
#endif
|
43
data_structures/kdtree/kdtree.hpp
Normal file
43
data_structures/kdtree/kdtree.hpp
Normal file
@ -0,0 +1,43 @@
|
|||||||
|
#ifndef GEOAPI_KDTREE_KDTREE_HPP
|
||||||
|
#define GEOAPI_KDTREE_KDTREE_HPP
|
||||||
|
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "build.hpp"
|
||||||
|
#include "nns.hpp"
|
||||||
|
|
||||||
|
namespace kd
|
||||||
|
{
|
||||||
|
|
||||||
|
template <class T, class U>
|
||||||
|
class KdTree
|
||||||
|
{
|
||||||
|
public:
|
||||||
|
KdTree() {}
|
||||||
|
|
||||||
|
template <class It>
|
||||||
|
KdTree(It first, It last);
|
||||||
|
|
||||||
|
const U& lookup(const Point<T>& pk) const;
|
||||||
|
|
||||||
|
protected:
|
||||||
|
std::unique_ptr<KdNode<float, U>> root;
|
||||||
|
};
|
||||||
|
|
||||||
|
template <class T, class U>
|
||||||
|
const U& KdTree<T, U>::lookup(const Point<T>& pk) const
|
||||||
|
{
|
||||||
|
// do a nearest neighbour search on the tree
|
||||||
|
return kd::nns(pk, root.get())->data;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <class T, class U>
|
||||||
|
template <class It>
|
||||||
|
KdTree<T, U>::KdTree(It first, It last)
|
||||||
|
{
|
||||||
|
root.reset(kd::build<T, U, It>(first, last));
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif
|
43
data_structures/kdtree/math.hpp
Normal file
43
data_structures/kdtree/math.hpp
Normal file
@ -0,0 +1,43 @@
|
|||||||
|
#ifndef GEOAPI_KDTREE_MATH_HPP
|
||||||
|
#define GEOAPI_KDTREE_MATH_HPP
|
||||||
|
|
||||||
|
#include <limits>
|
||||||
|
#include <cmath>
|
||||||
|
|
||||||
|
#include "point.hpp"
|
||||||
|
|
||||||
|
namespace kd {
|
||||||
|
namespace math {
|
||||||
|
|
||||||
|
using byte = unsigned char;
|
||||||
|
|
||||||
|
// returns the squared distance between two points
|
||||||
|
template<class T>
|
||||||
|
T distance_sq(const Point<T>& a, const Point<T>& b)
|
||||||
|
{
|
||||||
|
auto dx = a.longitude - b.longitude;
|
||||||
|
auto dy = a.latitude - b.latitude;
|
||||||
|
return dx * dx + dy * dy;
|
||||||
|
}
|
||||||
|
|
||||||
|
// returns the distance between two points
|
||||||
|
template<class T>
|
||||||
|
T distance(const Point<T>& a, const Point<T>& b)
|
||||||
|
{
|
||||||
|
return std::sqrt(distance_sq(a, b));
|
||||||
|
}
|
||||||
|
|
||||||
|
// returns the distance between two points looking at a specific axis
|
||||||
|
// \param axis 0 if abscissa else 1 if ordinate
|
||||||
|
template <class T>
|
||||||
|
T axial_distance(const Point<T>& a, const Point<T>& b, byte axis)
|
||||||
|
{
|
||||||
|
return axis == 0 ?
|
||||||
|
a.longitude - b.longitude:
|
||||||
|
a.latitude - b.latitude;
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif
|
104
data_structures/kdtree/nns.hpp
Normal file
104
data_structures/kdtree/nns.hpp
Normal file
@ -0,0 +1,104 @@
|
|||||||
|
#ifndef GEOAPI_KDTREE_NNS_HPP
|
||||||
|
#define GEOAPI_KDTREE_NNS_HPP
|
||||||
|
|
||||||
|
#include "math.hpp"
|
||||||
|
#include "point.hpp"
|
||||||
|
#include "kdnode.hpp"
|
||||||
|
|
||||||
|
namespace kd {
|
||||||
|
|
||||||
|
// helper class for calculating the nearest neighbour in a kdtree
|
||||||
|
template <class T, class U>
|
||||||
|
struct Result
|
||||||
|
{
|
||||||
|
Result()
|
||||||
|
: node(nullptr), distance_sq(std::numeric_limits<T>::infinity()) {}
|
||||||
|
|
||||||
|
Result(const KdNode<T, U>* node, T distance_sq)
|
||||||
|
: node(node), distance_sq(distance_sq) {}
|
||||||
|
|
||||||
|
const KdNode<T, U>* node;
|
||||||
|
T distance_sq;
|
||||||
|
};
|
||||||
|
|
||||||
|
|
||||||
|
// a recursive implementation for the kdtree nearest neighbour search
|
||||||
|
// \param p the point for which we search for the nearest neighbour
|
||||||
|
// \param node the root of the subtree during recursive descent
|
||||||
|
// \param best the place to save the best result so far
|
||||||
|
template <class T, class U>
|
||||||
|
void nns(const Point<T>& p, const KdNode<T, U>* const node, Result<T, U>& best)
|
||||||
|
{
|
||||||
|
if(node == nullptr)
|
||||||
|
return;
|
||||||
|
|
||||||
|
T d = math::distance_sq(p, node->coord);
|
||||||
|
|
||||||
|
// keep record of the closest point C found so far
|
||||||
|
if(d < best.distance_sq)
|
||||||
|
{
|
||||||
|
best.node = node;
|
||||||
|
best.distance_sq = d;
|
||||||
|
}
|
||||||
|
|
||||||
|
// where to traverse next?
|
||||||
|
// what to prune?
|
||||||
|
|
||||||
|
// |
|
||||||
|
// possible |
|
||||||
|
// prune *
|
||||||
|
// area | - - - - -* P
|
||||||
|
// |
|
||||||
|
//
|
||||||
|
// |----------|
|
||||||
|
// dx
|
||||||
|
//
|
||||||
|
|
||||||
|
// possible prune
|
||||||
|
// RIGHT area
|
||||||
|
//
|
||||||
|
// --------*------ ---
|
||||||
|
// | |
|
||||||
|
// LEFT |
|
||||||
|
// | | dy
|
||||||
|
// |
|
||||||
|
// | |
|
||||||
|
// * p ---
|
||||||
|
|
||||||
|
T axd = math::axial_distance(p, node->coord, node->axis);
|
||||||
|
|
||||||
|
// traverse the subtree in order that
|
||||||
|
// maximizes the probability for pruning
|
||||||
|
auto near = axd > 0 ? node->right : node->left;
|
||||||
|
auto far = axd > 0 ? node->left : node->right;
|
||||||
|
|
||||||
|
// try near first
|
||||||
|
nns(p, near, best);
|
||||||
|
|
||||||
|
// prune subtrees once their bounding boxes say
|
||||||
|
// that they can't contain any point closer than C
|
||||||
|
if(axd * axd >= best.distance_sq)
|
||||||
|
return;
|
||||||
|
|
||||||
|
// try other subtree
|
||||||
|
nns(p, far, best);
|
||||||
|
}
|
||||||
|
|
||||||
|
// an implementation for the kdtree nearest neighbour search
|
||||||
|
// \param p the point for which we search for the nearest neighbour
|
||||||
|
// \param root the root of the tree
|
||||||
|
// \return the nearest neighbour for the point p
|
||||||
|
template <class T, class U>
|
||||||
|
const KdNode<T, U>* nns(const Point<T>& p, const KdNode<T, U>* root)
|
||||||
|
{
|
||||||
|
Result<T, U> best;
|
||||||
|
|
||||||
|
// begin recursive search
|
||||||
|
nns(p, root, best);
|
||||||
|
|
||||||
|
return best.node;
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif
|
33
data_structures/kdtree/point.hpp
Normal file
33
data_structures/kdtree/point.hpp
Normal file
@ -0,0 +1,33 @@
|
|||||||
|
#ifndef GEOAPI_KDTREE_POINT_HPP
|
||||||
|
#define GEOAPI_KDTREE_POINT_HPP
|
||||||
|
|
||||||
|
#include <ostream>
|
||||||
|
|
||||||
|
namespace kd {
|
||||||
|
|
||||||
|
template <class T>
|
||||||
|
class Point
|
||||||
|
{
|
||||||
|
public:
|
||||||
|
Point(T latitude, T longitude)
|
||||||
|
: latitude(latitude), longitude(longitude) {}
|
||||||
|
|
||||||
|
// latitude
|
||||||
|
// y
|
||||||
|
// ^
|
||||||
|
// |
|
||||||
|
// 0---> x longitude
|
||||||
|
|
||||||
|
T latitude;
|
||||||
|
T longitude;
|
||||||
|
|
||||||
|
/// nice stream formatting with the standard << operator
|
||||||
|
friend std::ostream& operator<< (std::ostream& stream, const Point& p) {
|
||||||
|
return stream << "(lat: " << p.latitude
|
||||||
|
<< ", lng: " << p.longitude << ')';
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
#endif
|
Loading…
Reference in New Issue
Block a user