Add cg::kdtree::closer_than

This commit is contained in:
Nikita Lisitsa 2025-05-14 22:56:26 +03:00
parent 08ca6ab21b
commit a8e32e2d98

View file

@ -95,6 +95,12 @@ namespace psemek::cg
value_type const & closest(point_type const & target) const; value_type const & closest(point_type const & target) const;
template <typename Iterator>
Iterator closer_than(point_type const & target, scalar_type max_distance, Iterator out) const;
template <typename Callback>
void closer_than_map(point_type const & target, scalar_type max_distance, Callback && callback) const;
private: private:
using node_id = std::uint32_t; 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); 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; value_type const * closest_impl(point_type const & target, scalar_type & best_distance_sqr, node_id id, std::uint32_t split_axis) const;
template <typename Iterator>
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 <typename Callback>
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 <typename T, std::size_t N, typename Data> template <typename T, std::size_t N, typename Data>
@ -161,6 +173,26 @@ namespace psemek::cg
return *closest_impl(target, best_distance_sqr, 0, 0); return *closest_impl(target, best_distance_sqr, 0, 0);
} }
template <typename T, std::size_t N, typename Data>
template <typename Iterator>
Iterator kdtree<T, N, Data>::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 <typename T, std::size_t N, typename Data>
template <typename Callback>
void kdtree<T, N, Data>::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 <typename T, std::size_t N, typename Data> template <typename T, std::size_t N, typename Data>
template <typename Iterator> template <typename Iterator>
kdtree<T, N, Data>::node_id kdtree<T, N, Data>::build_node_impl(Iterator begin, Iterator end, std::uint32_t split_axis) kdtree<T, N, Data>::node_id kdtree<T, N, Data>::build_node_impl(Iterator begin, Iterator end, std::uint32_t split_axis)
@ -254,4 +286,46 @@ namespace psemek::cg
return result; return result;
} }
template <typename T, std::size_t N, typename Data>
template <typename Iterator>
Iterator kdtree<T, N, Data>::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 <typename T, std::size_t N, typename Data>
template <typename Callback>
void kdtree<T, N, Data>::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));
}
} }