From 98cc39ffd869ab604f85fe00119d034042226f74 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dominik=20Tomic=CC=8Cevic=CC=81?= Date: Fri, 19 Jun 2015 10:04:42 +0200 Subject: [PATCH] added kdtree implementation --- data_structures/kdtree/build.hpp | 70 ++++++++++++++++++++ data_structures/kdtree/kdnode.hpp | 48 ++++++++++++++ data_structures/kdtree/kdtree.hpp | 43 ++++++++++++ data_structures/kdtree/math.hpp | 43 ++++++++++++ data_structures/kdtree/nns.hpp | 104 ++++++++++++++++++++++++++++++ data_structures/kdtree/point.hpp | 33 ++++++++++ 6 files changed, 341 insertions(+) create mode 100644 data_structures/kdtree/build.hpp create mode 100644 data_structures/kdtree/kdnode.hpp create mode 100644 data_structures/kdtree/kdtree.hpp create mode 100644 data_structures/kdtree/math.hpp create mode 100644 data_structures/kdtree/nns.hpp create mode 100644 data_structures/kdtree/point.hpp diff --git a/data_structures/kdtree/build.hpp b/data_structures/kdtree/build.hpp new file mode 100644 index 000000000..dc0a857ae --- /dev/null +++ b/data_structures/kdtree/build.hpp @@ -0,0 +1,70 @@ +#ifndef MEMGRAPH_DATA_STRUCTURES_KDTREE_BUILD_HPP +#define MEMGRAPH_DATA_STRUCTURES_KDTREE_BUILD_HPP + +#include +#include +#include + +#include "math.hpp" +#include "kdnode.hpp" + +namespace kd { + +template +using Nodes = std::vector*>; + +template +KdNode* build(Nodes& 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* a, KdNode* 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(nodes.begin(), nodes.begin() + median); + auto right = Nodes(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 +KdNode* build(It first, It last) +{ + Nodes kdnodes; + + std::transform(first, last, std::back_inserter(kdnodes), + [&](const std::pair, U>& element) { + auto key = element.first; + auto data = element.second; + return new KdNode(key, data); + }); + + // build the tree from the kdnodes and return the root node + return build(kdnodes); +} + +} + +#endif diff --git a/data_structures/kdtree/kdnode.hpp b/data_structures/kdtree/kdnode.hpp new file mode 100644 index 000000000..b3a310aea --- /dev/null +++ b/data_structures/kdtree/kdnode.hpp @@ -0,0 +1,48 @@ +#ifndef GEOAPI_KDTREE_KDNODE_HPP +#define GEOAPI_KDTREE_KDNODE_HPP + +#include + +#include "point.hpp" + +namespace kd { + +template +class KdNode +{ +public: + KdNode(const U& data) + : axis(0), coord(Point(0, 0)), left(nullptr), right(nullptr), data(data) { } + + KdNode(const Point& coord, const U& data) + : axis(0), coord(coord), left(nullptr), right(nullptr), data(data) { } + + KdNode(unsigned char axis, const Point& coord, const U& data) + : axis(axis), coord(coord), left(nullptr), right(nullptr), data(data) { } + + KdNode(unsigned char axis, const Point& coord, KdNode* left, KdNode* right, const U& data) + : axis(axis), coord(coord), left(left), right(right), data(data) { } + + ~KdNode(); + + unsigned char axis; + + Point coord; + + KdNode* left; + KdNode* right; + + U data; +}; + +template +KdNode::~KdNode() +{ + delete left; + delete right; +} + +} + + +#endif diff --git a/data_structures/kdtree/kdtree.hpp b/data_structures/kdtree/kdtree.hpp new file mode 100644 index 000000000..9a1721dbc --- /dev/null +++ b/data_structures/kdtree/kdtree.hpp @@ -0,0 +1,43 @@ +#ifndef GEOAPI_KDTREE_KDTREE_HPP +#define GEOAPI_KDTREE_KDTREE_HPP + +#include + +#include "build.hpp" +#include "nns.hpp" + +namespace kd +{ + +template +class KdTree +{ +public: + KdTree() {} + + template + KdTree(It first, It last); + + const U& lookup(const Point& pk) const; + +protected: + std::unique_ptr> root; +}; + +template +const U& KdTree::lookup(const Point& pk) const +{ + // do a nearest neighbour search on the tree + return kd::nns(pk, root.get())->data; +} + +template +template +KdTree::KdTree(It first, It last) +{ + root.reset(kd::build(first, last)); +} + +} + +#endif diff --git a/data_structures/kdtree/math.hpp b/data_structures/kdtree/math.hpp new file mode 100644 index 000000000..52e5b5fb0 --- /dev/null +++ b/data_structures/kdtree/math.hpp @@ -0,0 +1,43 @@ +#ifndef GEOAPI_KDTREE_MATH_HPP +#define GEOAPI_KDTREE_MATH_HPP + +#include +#include + +#include "point.hpp" + +namespace kd { +namespace math { + +using byte = unsigned char; + +// returns the squared distance between two points +template +T distance_sq(const Point& a, const Point& 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 +T distance(const Point& a, const Point& 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 +T axial_distance(const Point& a, const Point& b, byte axis) +{ + return axis == 0 ? + a.longitude - b.longitude: + a.latitude - b.latitude; +} + +} +} + +#endif diff --git a/data_structures/kdtree/nns.hpp b/data_structures/kdtree/nns.hpp new file mode 100644 index 000000000..cc0cb2820 --- /dev/null +++ b/data_structures/kdtree/nns.hpp @@ -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 +struct Result +{ + Result() + : node(nullptr), distance_sq(std::numeric_limits::infinity()) {} + + Result(const KdNode* node, T distance_sq) + : node(node), distance_sq(distance_sq) {} + + const KdNode* 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 +void nns(const Point& p, const KdNode* const node, Result& 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 +const KdNode* nns(const Point& p, const KdNode* root) +{ + Result best; + + // begin recursive search + nns(p, root, best); + + return best.node; +} + +} + +#endif diff --git a/data_structures/kdtree/point.hpp b/data_structures/kdtree/point.hpp new file mode 100644 index 000000000..26c03fd27 --- /dev/null +++ b/data_structures/kdtree/point.hpp @@ -0,0 +1,33 @@ +#ifndef GEOAPI_KDTREE_POINT_HPP +#define GEOAPI_KDTREE_POINT_HPP + +#include + +namespace kd { + +template +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