diff --git a/libs/cg/include/psemek/cg/kdtree.hpp b/libs/cg/include/psemek/cg/kdtree.hpp index 961531fb..93666fa6 100644 --- a/libs/cg/include/psemek/cg/kdtree.hpp +++ b/libs/cg/include/psemek/cg/kdtree.hpp @@ -95,6 +95,12 @@ namespace psemek::cg value_type const & closest(point_type const & target) const; + template + Iterator closer_than(point_type const & target, scalar_type max_distance, Iterator out) const; + + template + void closer_than_map(point_type const & target, scalar_type max_distance, Callback && callback) const; + private: using node_id = std::uint32_t; @@ -122,6 +128,12 @@ namespace psemek::cg bool insert_impl(value_type && value, node_id id, std::uint32_t split_axis); value_type const * closest_impl(point_type const & target, scalar_type & best_distance_sqr, node_id id, std::uint32_t split_axis) const; + + template + Iterator closer_than_impl(point_type const & target, scalar_type max_distance_sqr, Iterator out, node_id id, std::uint32_t split_axis) const; + + template + void closer_than_map_impl(point_type const & target, scalar_type max_distance_sqr, Callback & callback, node_id id, std::uint32_t split_axis) const; }; template @@ -161,6 +173,26 @@ namespace psemek::cg return *closest_impl(target, best_distance_sqr, 0, 0); } + template + template + Iterator kdtree::closer_than(point_type const & target, scalar_type max_distance, Iterator out) const + { + if (nodes_.empty()) + return out; + + return closer_than_impl(target, math::sqr(max_distance), out, 0, 0); + } + + template + template + void kdtree::closer_than_map(point_type const & target, scalar_type max_distance, Callback && callback) const + { + if (nodes_.empty()) + return; + + return closer_than_map_impl(target, math::sqr(max_distance), callback, 0, 0); + } + template template kdtree::node_id kdtree::build_node_impl(Iterator begin, Iterator end, std::uint32_t split_axis) @@ -254,4 +286,46 @@ namespace psemek::cg return result; } + template + template + Iterator kdtree::closer_than_impl(point_type const & target, scalar_type max_distance_sqr, Iterator out, node_id id, std::uint32_t split_axis) const + { + auto const & node = nodes_[id]; + auto const & node_point = detail::get_point(node.value); + + if (math::distance_sqr(node_point, target) < max_distance_sqr) + *out++ = node.value; + + auto delta = target[split_axis] - node_point[split_axis]; + auto delta_sqr = math::sqr(delta); + + if (node.children[0] != null && (delta < 0 || delta_sqr < max_distance_sqr)) + out = closer_than_impl(target, max_distance_sqr, out, node.children[0], next_axis(split_axis)); + + if (node.children[1] != null && (delta >= 0 || delta_sqr < max_distance_sqr)) + out = closer_than_impl(target, max_distance_sqr, out, node.children[1], next_axis(split_axis)); + + return out; + } + + template + template + void kdtree::closer_than_map_impl(point_type const & target, scalar_type max_distance_sqr, Callback & callback, node_id id, std::uint32_t split_axis) const + { + auto const & node = nodes_[id]; + auto const & node_point = detail::get_point(node.value); + + if (math::distance_sqr(node_point, target) < max_distance_sqr) + callback(node.value); + + auto delta = target[split_axis] - node_point[split_axis]; + auto delta_sqr = math::sqr(delta); + + if (node.children[0] != null && (delta < 0 || delta_sqr < max_distance_sqr)) + closer_than_map_impl(target, max_distance_sqr, callback, node.children[0], next_axis(split_axis)); + + if (node.children[1] != null && (delta >= 0 || delta_sqr < max_distance_sqr)) + closer_than_map_impl(target, max_distance_sqr, callback, node.children[1], next_axis(split_axis)); + } + }