|
|
@@ -0,0 +1,441 @@
|
|
|
+
|
|
|
+// Copyright 2018 Baidu. Inc All Rights Reserved
|
|
|
+// Author: Fuxiangyu (fuxiangyu@baidu.com)
|
|
|
+// Date: 2018-06-19 14:44
|
|
|
+//
|
|
|
+// Description: aaboxkdtree for 3d element
|
|
|
+
|
|
|
+#ifndef _MATH_COMMON_AABOXKDTREE3D_H_
|
|
|
+#define _MATH_COMMON_AABOXKDTREE3D_H_
|
|
|
+
|
|
|
+#pragma once
|
|
|
+
|
|
|
+#include "math/aabox3d.h"
|
|
|
+#include "math/math_utils.h"
|
|
|
+
|
|
|
+#include <algorithm>
|
|
|
+#include <limits>
|
|
|
+#include <memory>
|
|
|
+#include <vector>
|
|
|
+
|
|
|
+namespace decision::math {
|
|
|
+
|
|
|
+struct AABox3dKDTreeParams {
|
|
|
+ // The maximum depth of the kdtree.
|
|
|
+ int max_depth = -1;
|
|
|
+ // The maximum number of items in one leaf node.
|
|
|
+ int max_leaf_size = -1;
|
|
|
+ // The maximum dimension size of leaf node.
|
|
|
+ double max_leaf_dimension = -1.0;
|
|
|
+};
|
|
|
+
|
|
|
+template <class ObjectType>
|
|
|
+class AABoxKDTree3dNode {
|
|
|
+public:
|
|
|
+ using ObjectPtr = const ObjectType*;
|
|
|
+
|
|
|
+ AABoxKDTree3dNode(const std::vector<ObjectPtr>& objects,
|
|
|
+ const AABox3dKDTreeParams& params, int depth)
|
|
|
+ : _depth(depth)
|
|
|
+ {
|
|
|
+ CHECK(!objects.empty());
|
|
|
+
|
|
|
+ compute_boundary(objects);
|
|
|
+ compute_partition();
|
|
|
+
|
|
|
+ if (split_to_subnodes(objects, params)) {
|
|
|
+ std::vector<ObjectPtr> left_subnode_objects;
|
|
|
+ std::vector<ObjectPtr> right_subnode_objects;
|
|
|
+ partition_objects(objects, &left_subnode_objects, &right_subnode_objects);
|
|
|
+
|
|
|
+ // Split to sub-nodes.
|
|
|
+ if (!left_subnode_objects.empty()) {
|
|
|
+ _left_subnode.reset(new AABoxKDTree3dNode<ObjectType>(
|
|
|
+ left_subnode_objects, params, depth + 1));
|
|
|
+ }
|
|
|
+ if (!right_subnode_objects.empty()) {
|
|
|
+ _right_subnode.reset(new AABoxKDTree3dNode<ObjectType>(
|
|
|
+ right_subnode_objects, params, depth + 1));
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ init_objects(objects);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ ObjectPtr get_nearest_object(const Vec3d& point) const
|
|
|
+ {
|
|
|
+ ObjectPtr nearest_object = nullptr;
|
|
|
+ double min_distance_sqr = std::numeric_limits<double>::infinity();
|
|
|
+ get_nearest_object_internal(point, &min_distance_sqr, &nearest_object);
|
|
|
+ return nearest_object;
|
|
|
+ }
|
|
|
+
|
|
|
+ std::vector<ObjectPtr> get_objects(const Vec3d& point, const double distance) const
|
|
|
+ {
|
|
|
+ std::vector<ObjectPtr> result_objects;
|
|
|
+ get_objects_internal(point, distance, sqr(distance), &result_objects);
|
|
|
+ return result_objects;
|
|
|
+ }
|
|
|
+
|
|
|
+ AABox3d get_bounding_box() const
|
|
|
+ {
|
|
|
+ return AABox3d({ _min_x, _min_y, _min_z }, { _max_x, _max_y, _max_z });
|
|
|
+ }
|
|
|
+
|
|
|
+private:
|
|
|
+ void init_objects(const std::vector<ObjectPtr>& objects)
|
|
|
+ {
|
|
|
+ _num_objects = objects.size();
|
|
|
+ _objects_sorted_by_min = objects;
|
|
|
+ _objects_sorted_by_max = objects;
|
|
|
+ std::sort(_objects_sorted_by_min.begin(), _objects_sorted_by_min.end(),
|
|
|
+ [&](ObjectPtr obj1, ObjectPtr obj2) {
|
|
|
+ return _partition == PARTITION_X ? obj1->aabox().min_x() < obj2->aabox().min_x()
|
|
|
+ : (_partition == PARTITION_Y ? obj1->aabox().min_y() < obj2->aabox().min_y()
|
|
|
+ : obj1->aabox().min_z() < obj2->aabox().min_z());
|
|
|
+ });
|
|
|
+ std::sort(_objects_sorted_by_max.begin(), _objects_sorted_by_max.end(),
|
|
|
+ [&](ObjectPtr obj1, ObjectPtr obj2) {
|
|
|
+ return _partition == PARTITION_X ? obj1->aabox().max_x() > obj2->aabox().max_x()
|
|
|
+ : (_partition == PARTITION_Y ? obj1->aabox().max_y() > obj2->aabox().max_y()
|
|
|
+ : obj1->aabox().max_z() > obj2->aabox().max_z());
|
|
|
+ });
|
|
|
+ _objects_sorted_by_min_bound.reserve(_num_objects);
|
|
|
+ for (ObjectPtr object : _objects_sorted_by_min) {
|
|
|
+ _objects_sorted_by_min_bound.push_back(
|
|
|
+ _partition == PARTITION_X ? object->aabox().min_x()
|
|
|
+ : (_partition == PARTITION_Y ? object->aabox().min_y()
|
|
|
+ : object->aabox().min_z()));
|
|
|
+ }
|
|
|
+ _objects_sorted_by_max_bound.reserve(_num_objects);
|
|
|
+ for (ObjectPtr object : _objects_sorted_by_max) {
|
|
|
+ _objects_sorted_by_max_bound.push_back(
|
|
|
+ _partition == PARTITION_X ? object->aabox().max_x()
|
|
|
+ : (_partition == PARTITION_Y ? object->aabox().max_y()
|
|
|
+ : object->aabox().max_z()));
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ bool split_to_subnodes(const std::vector<ObjectPtr>& objects, const AABox3dKDTreeParams& params)
|
|
|
+ {
|
|
|
+ if (params.max_depth >= 0 && _depth >= params.max_depth) {
|
|
|
+ return false;
|
|
|
+ }
|
|
|
+ if (static_cast<int>(objects.size()) <= std::max(1, params.max_leaf_size)) {
|
|
|
+ return false;
|
|
|
+ }
|
|
|
+ double max_dim = std::max(_max_x - _min_x, _max_y - _min_y);
|
|
|
+ max_dim = std::max(max_dim, _max_z - _min_z);
|
|
|
+ if (params.max_leaf_dimension >= 0.0 && max_dim <= params.max_leaf_dimension) {
|
|
|
+ return false;
|
|
|
+ }
|
|
|
+ return true;
|
|
|
+ }
|
|
|
+
|
|
|
+ double lowerbound_distance_sqr_to_point(const Vec3d& point) const
|
|
|
+ {
|
|
|
+ double dx = 0.0;
|
|
|
+ if (point.x() < _min_x) {
|
|
|
+ dx = _min_x - point.x();
|
|
|
+ } else if (point.x() > _max_x) {
|
|
|
+ dx = point.x() - _max_x;
|
|
|
+ }
|
|
|
+ double dy = 0.0;
|
|
|
+ if (point.y() < _min_y) {
|
|
|
+ dy = _min_y - point.y();
|
|
|
+ } else if (point.y() > _max_y) {
|
|
|
+ dy = point.y() - _max_y;
|
|
|
+ }
|
|
|
+ double dz = 0.0;
|
|
|
+ if (point.z() < _min_z) {
|
|
|
+ dz = _min_z - point.z();
|
|
|
+ } else if (point.z() > _max_z) {
|
|
|
+ dz = point.z() - _max_z;
|
|
|
+ }
|
|
|
+ return dx * dx + dy * dy + dz * dz;
|
|
|
+ }
|
|
|
+
|
|
|
+ double upperbound_distance_sqr_to_point(const Vec3d& point) const
|
|
|
+ {
|
|
|
+ const double dx = (point.x() > _mid_x ? (point.x() - _min_x) : (point.x() - _max_x));
|
|
|
+ const double dy = (point.y() > _mid_y ? (point.y() - _min_y) : (point.y() - _max_y));
|
|
|
+ const double dz = (point.z() > _mid_z ? (point.z() - _min_z) : (point.z() - _max_z));
|
|
|
+ return dx * dx + dy * dy + dz * dz;
|
|
|
+ }
|
|
|
+
|
|
|
+ void get_all_objects(std::vector<ObjectPtr>* const result_objects) const
|
|
|
+ {
|
|
|
+ result_objects->insert(result_objects->end(),
|
|
|
+ _objects_sorted_by_min.begin(), _objects_sorted_by_min.end());
|
|
|
+ if (_left_subnode != nullptr) {
|
|
|
+ _left_subnode->get_all_objects(result_objects);
|
|
|
+ }
|
|
|
+ if (_right_subnode != nullptr) {
|
|
|
+ _right_subnode->get_all_objects(result_objects);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ void get_objects_internal(const Vec3d& point,
|
|
|
+ const double distance,
|
|
|
+ const double distance_sqr,
|
|
|
+ std::vector<ObjectPtr>* const result_objects) const
|
|
|
+ {
|
|
|
+ if (lowerbound_distance_sqr_to_point(point) > distance_sqr) {
|
|
|
+ return;
|
|
|
+ }
|
|
|
+ if (upperbound_distance_sqr_to_point(point) <= distance_sqr) {
|
|
|
+ get_all_objects(result_objects);
|
|
|
+ return;
|
|
|
+ }
|
|
|
+ const double pvalue = (_partition == PARTITION_X ? point.x()
|
|
|
+ : (_partition == PARTITION_Y ? point.y() : point.z()));
|
|
|
+ if (pvalue < _partition_position) {
|
|
|
+ const double limit = pvalue + distance;
|
|
|
+ for (int i = 0; i < _num_objects; ++i) {
|
|
|
+ if (_objects_sorted_by_min_bound[i] > limit) {
|
|
|
+ break;
|
|
|
+ }
|
|
|
+ ObjectPtr object = _objects_sorted_by_min[i];
|
|
|
+ if (object->distance_sqr_to(point) <= distance_sqr) {
|
|
|
+ result_objects->push_back(object);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ const double limit = pvalue - distance;
|
|
|
+ for (int i = 0; i < _num_objects; ++i) {
|
|
|
+ if (_objects_sorted_by_max_bound[i] < limit) {
|
|
|
+ break;
|
|
|
+ }
|
|
|
+ ObjectPtr object = _objects_sorted_by_max[i];
|
|
|
+ if (object->distance_sqr_to(point) <= distance_sqr) {
|
|
|
+ result_objects->push_back(object);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ if (_left_subnode != nullptr) {
|
|
|
+ _left_subnode->get_objects_internal(point, distance, distance_sqr, result_objects);
|
|
|
+ }
|
|
|
+ if (_right_subnode != nullptr) {
|
|
|
+ _right_subnode->get_objects_internal(point, distance, distance_sqr, result_objects);
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ void get_nearest_object_internal(const Vec3d& point,
|
|
|
+ double* const min_distance_sqr,
|
|
|
+ ObjectPtr* const nearest_object) const
|
|
|
+ {
|
|
|
+ if (lowerbound_distance_sqr_to_point(point) >= *min_distance_sqr - kMathEpsilon) {
|
|
|
+ return;
|
|
|
+ }
|
|
|
+ const double pvalue = (_partition == PARTITION_X ? point.x()
|
|
|
+ : (_partition == PARTITION_Y ? point.y() : point.z()));
|
|
|
+ const bool search_left_first = (pvalue < _partition_position);
|
|
|
+ if (search_left_first) {
|
|
|
+ if (_left_subnode != nullptr) {
|
|
|
+ _left_subnode->get_nearest_object_internal(
|
|
|
+ point, min_distance_sqr, nearest_object);
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ if (_right_subnode != nullptr) {
|
|
|
+ _right_subnode->get_nearest_object_internal(
|
|
|
+ point, min_distance_sqr, nearest_object);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ if (*min_distance_sqr <= kMathEpsilon) {
|
|
|
+ return;
|
|
|
+ }
|
|
|
+
|
|
|
+ if (search_left_first) {
|
|
|
+ for (int i = 0; i < _num_objects; ++i) {
|
|
|
+ const double bound = _objects_sorted_by_min_bound[i];
|
|
|
+ if (bound > pvalue && sqr(bound - pvalue) > *min_distance_sqr) {
|
|
|
+ break;
|
|
|
+ }
|
|
|
+ ObjectPtr object = _objects_sorted_by_min[i];
|
|
|
+ const double distance_sqr = object->distance_sqr_to(point);
|
|
|
+ if (distance_sqr < *min_distance_sqr) {
|
|
|
+ *min_distance_sqr = distance_sqr;
|
|
|
+ *nearest_object = object;
|
|
|
+ }
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ for (int i = 0; i < _num_objects; ++i) {
|
|
|
+ const double bound = _objects_sorted_by_max_bound[i];
|
|
|
+ if (bound < pvalue && sqr(bound - pvalue) > *min_distance_sqr) {
|
|
|
+ break;
|
|
|
+ }
|
|
|
+ ObjectPtr object = _objects_sorted_by_max[i];
|
|
|
+ const double distance_sqr = object->distance_sqr_to(point);
|
|
|
+ if (distance_sqr < *min_distance_sqr) {
|
|
|
+ *min_distance_sqr = distance_sqr;
|
|
|
+ *nearest_object = object;
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ if (*min_distance_sqr <= kMathEpsilon) {
|
|
|
+ return;
|
|
|
+ }
|
|
|
+ if (search_left_first) {
|
|
|
+ if (_right_subnode != nullptr) {
|
|
|
+ _right_subnode->get_nearest_object_internal(
|
|
|
+ point, min_distance_sqr, nearest_object);
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ if (_left_subnode != nullptr) {
|
|
|
+ _left_subnode->get_nearest_object_internal(
|
|
|
+ point, min_distance_sqr, nearest_object);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ void compute_boundary(const std::vector<ObjectPtr>& objects)
|
|
|
+ {
|
|
|
+ _min_x = std::numeric_limits<double>::infinity();
|
|
|
+ _min_y = std::numeric_limits<double>::infinity();
|
|
|
+ _min_z = std::numeric_limits<double>::infinity();
|
|
|
+ _max_x = -std::numeric_limits<double>::infinity();
|
|
|
+ _max_y = -std::numeric_limits<double>::infinity();
|
|
|
+ _max_z = -std::numeric_limits<double>::infinity();
|
|
|
+ for (ObjectPtr object : objects) {
|
|
|
+ _min_x = std::min(_min_x, object->aabox().min_x());
|
|
|
+ _max_x = std::max(_max_x, object->aabox().max_x());
|
|
|
+ _min_y = std::min(_min_y, object->aabox().min_y());
|
|
|
+ _max_y = std::max(_max_y, object->aabox().max_y());
|
|
|
+ _min_z = std::min(_min_z, object->aabox().min_z());
|
|
|
+ _max_z = std::max(_max_z, object->aabox().max_z());
|
|
|
+ }
|
|
|
+ _mid_x = (_min_x + _max_x) / 2.0;
|
|
|
+ _mid_y = (_min_y + _max_y) / 2.0;
|
|
|
+ _mid_z = (_min_z + _max_z) / 2.0;
|
|
|
+ }
|
|
|
+
|
|
|
+ void compute_partition()
|
|
|
+ {
|
|
|
+ double dx = _max_x - _min_x;
|
|
|
+ double dy = _max_y - _min_y;
|
|
|
+ double dz = _max_z - _min_z;
|
|
|
+ if (dx >= dy && dx >= dz) {
|
|
|
+ _partition = PARTITION_X;
|
|
|
+ _partition_position = (_min_x + _max_x) / 2.0;
|
|
|
+ } else if (dy >= dx && dy >= dz) {
|
|
|
+ _partition = PARTITION_Y;
|
|
|
+ _partition_position = (_min_y + _max_y) / 2.0;
|
|
|
+ } else {
|
|
|
+ _partition = PARTITION_Z;
|
|
|
+ _partition_position = (_min_z + _max_z) / 2.0;
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ void partition_objects(const std::vector<ObjectPtr>& objects,
|
|
|
+ std::vector<ObjectPtr>* const left_subnode_objects,
|
|
|
+ std::vector<ObjectPtr>* const right_subnode_objects)
|
|
|
+ {
|
|
|
+ left_subnode_objects->clear();
|
|
|
+ right_subnode_objects->clear();
|
|
|
+ std::vector<ObjectPtr> other_objects;
|
|
|
+ if (_partition == PARTITION_X) {
|
|
|
+ for (ObjectPtr object : objects) {
|
|
|
+ if (object->aabox().max_x() <= _partition_position) {
|
|
|
+ left_subnode_objects->push_back(object);
|
|
|
+ } else if (object->aabox().min_x() >= _partition_position) {
|
|
|
+ right_subnode_objects->push_back(object);
|
|
|
+ } else {
|
|
|
+ other_objects.push_back(object);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ } else if (_partition == PARTITION_Y) {
|
|
|
+ for (ObjectPtr object : objects) {
|
|
|
+ if (object->aabox().max_y() <= _partition_position) {
|
|
|
+ left_subnode_objects->push_back(object);
|
|
|
+ } else if (object->aabox().min_y() >= _partition_position) {
|
|
|
+ right_subnode_objects->push_back(object);
|
|
|
+ } else {
|
|
|
+ other_objects.push_back(object);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ } else {
|
|
|
+ for (ObjectPtr object : objects) {
|
|
|
+ if (object->aabox().max_z() <= _partition_position) {
|
|
|
+ left_subnode_objects->push_back(object);
|
|
|
+ } else if (object->aabox().min_z() >= _partition_position) {
|
|
|
+ right_subnode_objects->push_back(object);
|
|
|
+ } else {
|
|
|
+ other_objects.push_back(object);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ init_objects(other_objects);
|
|
|
+ }
|
|
|
+
|
|
|
+private:
|
|
|
+ int _num_objects = 0;
|
|
|
+ std::vector<ObjectPtr> _objects_sorted_by_min;
|
|
|
+ std::vector<ObjectPtr> _objects_sorted_by_max;
|
|
|
+ std::vector<double> _objects_sorted_by_min_bound;
|
|
|
+ std::vector<double> _objects_sorted_by_max_bound;
|
|
|
+ int _depth = 0;
|
|
|
+
|
|
|
+ // Boundary
|
|
|
+ double _min_x = 0.0;
|
|
|
+ double _max_x = 0.0;
|
|
|
+ double _min_y = 0.0;
|
|
|
+ double _max_y = 0.0;
|
|
|
+ double _min_z = 0.0;
|
|
|
+ double _max_z = 0.0;
|
|
|
+ double _mid_x = 0.0;
|
|
|
+ double _mid_y = 0.0;
|
|
|
+ double _mid_z = 0.0;
|
|
|
+
|
|
|
+ enum Partition {
|
|
|
+ PARTITION_X = 1,
|
|
|
+ PARTITION_Y = 2,
|
|
|
+ PARTITION_Z = 3,
|
|
|
+ };
|
|
|
+ Partition _partition = PARTITION_X;
|
|
|
+ double _partition_position = 0.0;
|
|
|
+
|
|
|
+ std::unique_ptr<AABoxKDTree3dNode<ObjectType>> _left_subnode = nullptr;
|
|
|
+ std::unique_ptr<AABoxKDTree3dNode<ObjectType>> _right_subnode = nullptr;
|
|
|
+};
|
|
|
+
|
|
|
+template <class ObjectType>
|
|
|
+class AABoxKDTree3d {
|
|
|
+public:
|
|
|
+ using ObjectPtr = const ObjectType*;
|
|
|
+
|
|
|
+ AABoxKDTree3d(const std::vector<ObjectType>& objects, const AABox3dKDTreeParams& params)
|
|
|
+ {
|
|
|
+ if (!objects.empty()) {
|
|
|
+ std::vector<ObjectPtr> object_ptrs;
|
|
|
+ for (const auto& object : objects) {
|
|
|
+ object_ptrs.push_back(&object);
|
|
|
+ }
|
|
|
+ _root.reset(new AABoxKDTree3dNode<ObjectType>(object_ptrs, params, 0));
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ ObjectPtr get_nearest_object(const Vec3d& point) const
|
|
|
+ {
|
|
|
+ return _root == nullptr ? nullptr : _root->get_nearest_object(point);
|
|
|
+ }
|
|
|
+
|
|
|
+ std::vector<ObjectPtr> get_objects(const Vec3d& point, const double distance) const
|
|
|
+ {
|
|
|
+ if (_root == nullptr) {
|
|
|
+ return {};
|
|
|
+ }
|
|
|
+ return _root->get_objects(point, distance);
|
|
|
+ }
|
|
|
+
|
|
|
+ AABox3d get_bounding_box() const
|
|
|
+ {
|
|
|
+ return _root == nullptr ? AABox3d() : _root->get_bounding_box();
|
|
|
+ }
|
|
|
+
|
|
|
+private:
|
|
|
+ std::unique_ptr<AABoxKDTree3dNode<ObjectType>> _root = nullptr;
|
|
|
+};
|
|
|
+
|
|
|
+}
|
|
|
+
|
|
|
+#endif // _MATH_COMMON_AABOXKDTREE3D_H_
|