kd_tree.h

Go to the documentation of this file.

/*
 *  Copyright (c) 2000-2022 Inria
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions are met:
 *
 *  * Redistributions of source code must retain the above copyright notice,
 *  this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright notice,
 *  this list of conditions and the following disclaimer in the documentation
 *  and/or other materials provided with the distribution.
 *  * Neither the name of the ALICE Project-Team nor the names of its
 *  contributors may be used to endorse or promote products derived from this
 *  software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 *  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 *  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 *  ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 *  LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 *  CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 *  SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 *  INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 *  CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 *  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *
 *  Contact: Bruno Levy
 *
 *     https://www.inria.fr/fr/bruno-levy
 *
 *     Inria,
 *     Domaine de Voluceau,
 *     78150 Le Chesnay - Rocquencourt
 *     FRANCE
 *
 */
 
#ifndef GEOGRAM_POINTS_KD_TREE
#define GEOGRAM_POINTS_KD_TREE
 
#include <geogram/basic/common.h>
#include <geogram/points/nn_search.h>
#include <algorithm>
 
namespace GEO {
 
    class GEOGRAM_API KdTree : public NearestNeighborSearch {
    public:
        KdTree(coord_index_t dim);
 
        void set_points(index_t nb_points, const double* points) override;
 
        bool stride_supported() const override;
 
        void set_points(
            index_t nb_points, const double* points, index_t stride
        ) override;
 
        void get_nearest_neighbors(
            index_t nb_neighbors,
            const double* query_point,
            index_t* neighbors,
            double* neighbors_sq_dist
        ) const override;
 
        void get_nearest_neighbors(
            index_t nb_neighbors,
            const double* query_point,
            index_t* neighbors,
            double* neighbors_sq_dist,
            KeepInitialValues
        ) const override;
 
        void get_nearest_neighbors(
            index_t nb_neighbors,
            index_t query_point,
            index_t* neighbors,
            double* neighbors_sq_dist
        ) const override;
 
        /**********************************************************************/
 
        struct NearestNeighbors {
 
            NearestNeighbors(
                index_t nb_neighbors_in,
                index_t* user_neighbors_in,
                double* user_neighbors_sq_dist_in,
                index_t* work_neighbors_in,
                double* work_neighbors_sq_dist_in
            ) :
                nb_neighbors(0),
                nb_neighbors_max(nb_neighbors_in),
                neighbors(work_neighbors_in),
                neighbors_sq_dist(work_neighbors_sq_dist_in),
                user_neighbors(user_neighbors_in),
                user_neighbors_sq_dist(user_neighbors_sq_dist_in),
                nb_visited(0)
                {
                    // Yes, '<=' because we got space for n+1 neigbors
                    // in the work arrays.
                    for(index_t i = 0; i <= nb_neighbors; ++i) {
                        neighbors[i] = index_t(-1);
                        neighbors_sq_dist[i] = Numeric::max_float64();
                    }
                }
 
            double furthest_neighbor_sq_dist() const {
                return
                    nb_neighbors == nb_neighbors_max ?
                    neighbors_sq_dist[nb_neighbors - 1] :
                    Numeric::max_float64()
                    ;
            }
 
            void insert(
                index_t neighbor, double sq_dist
            ) {
                geo_debug_assert(
                    sq_dist <= furthest_neighbor_sq_dist()
                );
 
                int i;
                for(i=int(nb_neighbors); i>0; --i) {
                    if(neighbors_sq_dist[i - 1] < sq_dist) {
                        break;
                    }
                    neighbors[i] = neighbors[i - 1];
                    neighbors_sq_dist[i] = neighbors_sq_dist[i - 1];
                }
 
                neighbors[i] = neighbor;
                neighbors_sq_dist[i] = sq_dist;
 
                if(nb_neighbors < nb_neighbors_max) {
                    ++nb_neighbors;
                }
            }
 
            void copy_from_user() {
                for(index_t i=0; i<nb_neighbors_max; ++i) {
                    neighbors[i] = user_neighbors[i];
                    neighbors_sq_dist[i] = user_neighbors_sq_dist[i];
                }
                neighbors[nb_neighbors_max] = index_t(-1);
                neighbors_sq_dist[nb_neighbors_max] = Numeric::max_float64();
                nb_neighbors = nb_neighbors_max;
            }
 
            void copy_to_user() {
                for(index_t i=0; i<nb_neighbors_max; ++i) {
                    user_neighbors[i] = neighbors[i];
                    user_neighbors_sq_dist[i] = neighbors_sq_dist[i];
                }
            }
 
            index_t nb_neighbors;
 
            index_t nb_neighbors_max;
 
            index_t* neighbors;
 
            double* neighbors_sq_dist;
 
            index_t* user_neighbors;
 
            double* user_neighbors_sq_dist;
 
            size_t nb_visited;
        };
 
        virtual void get_nearest_neighbors_recursive(
            index_t node_index, index_t b, index_t e,
            double* bbox_min, double* bbox_max,
            double bbox_dist, const double* query_point,
            NearestNeighbors& neighbors
        ) const;
 
        void init_bbox_and_bbox_dist_for_traversal(
            double* bbox_min, double* bbox_max,
            double& box_dist, const double* query_point
        ) const;
 
        index_t root() const {
            return root_;
        }
 
    protected:
        static constexpr index_t MAX_LEAF_SIZE = 16;
 
        virtual index_t build_tree() = 0 ;
 
        virtual void get_node(
            index_t n, index_t b, index_t e,
            index_t& left_child, index_t& right_child,
            coord_index_t&  splitting_coord,
            index_t& m,
            double& splitting_val
        ) const = 0;
 
 
 
        virtual void get_nearest_neighbors_leaf(
            index_t node_index, index_t b, index_t e,
            const double* query_point,
            NearestNeighbors& neighbors
        ) const;
 
        void get_minmax(
            index_t b, index_t e, coord_index_t coord,
            double& minval, double& maxval
        ) const {
            minval = Numeric::max_float64();
            maxval = Numeric::min_float64();
            for(index_t i = b; i < e; ++i) {
                double val = point_ptr(point_index_[i])[coord];
                minval = std::min(minval, val);
                maxval = std::max(maxval, val);
            }
        }
 
        double spread(index_t b, index_t e, coord_index_t coord) const {
            double minval,maxval;
            get_minmax(b,e,coord,minval,maxval);
            return maxval - minval;
        }
 
        ~KdTree() override;
 
    protected:
        vector<index_t> point_index_;
        vector<double> bbox_min_;
        vector<double> bbox_max_;
        index_t root_;
    };
 
    /*********************************************************************/
 
    class GEOGRAM_API BalancedKdTree : public KdTree {
    public:
        BalancedKdTree(coord_index_t dim);
 
    protected:
        ~BalancedKdTree() override;
 
        static index_t max_node_index(
            index_t node_id, index_t b, index_t e
        ) {
            if(e - b <= MAX_LEAF_SIZE) {
                return node_id;
            }
            index_t m = b + (e - b) / 2;
            return std::max(
                max_node_index(2 * node_id, b, m),
                max_node_index(2 * node_id + 1, m, e)
            );
        }
 
        coord_index_t best_splitting_coord(index_t b, index_t e);
 
        void create_kd_tree_recursive(
            index_t node_index, index_t b, index_t e
        ) {
            if(e - b <= MAX_LEAF_SIZE) {
                return;
            }
            index_t m = split_kd_node(node_index, b, e);
            create_kd_tree_recursive(2 * node_index, b, m);
            create_kd_tree_recursive(2 * node_index + 1, m, e);
        }
 
        index_t split_kd_node(
            index_t node_index, index_t b, index_t e
        );
 
        index_t build_tree() override;
 
        void get_node(
            index_t n, index_t b, index_t e,
            index_t& left_child, index_t& right_child,
            coord_index_t&  splitting_coord,
            index_t& m,
            double& splitting_val
        ) const override;
 
    protected:
 
        vector<coord_index_t> splitting_coord_;
 
        vector<double> splitting_val_;
 
        index_t m0_, m1_, m2_, m3_, m4_, m5_, m6_, m7_, m8_;
    };
 
    /*********************************************************************/
 
    class GEOGRAM_API AdaptiveKdTree : public KdTree {
    public:
        AdaptiveKdTree(coord_index_t dim);
 
    protected:
        index_t build_tree() override;
 
        void get_node(
            index_t n, index_t b, index_t e,
            index_t& left_child, index_t& right_child,
            coord_index_t&  splitting_coord,
            index_t& m,
            double& splitting_val
        ) const override;
 
        virtual index_t create_kd_tree_recursive(
            index_t b, index_t e,
            double* bbox_min, double* bbox_max
        );
 
        virtual void split_kd_node(
            index_t b, index_t e,
            double* bbox_min, double* bbox_max,
            index_t& m, coord_index_t& cut_dim, double& cut_val
        );
 
        virtual void plane_split(
            index_t b, index_t e, coord_index_t coord, double val,
            index_t& br1, index_t& br2
        );
 
        double point_coord(int index, coord_index_t coord) {
            geo_debug_assert(index >= 0);
            geo_debug_assert(index_t(index) < nb_points());
            geo_debug_assert(coord < dimension());
            index_t direct_index = point_index_[index_t(index)];
            geo_debug_assert(direct_index < nb_points());
            return (points_ + direct_index * stride_)[coord];
        }
 
 
        index_t nb_nodes() const {
            return splitting_coord_.size();
        }
 
        virtual index_t new_node();
 
    protected:
        vector<coord_index_t> splitting_coord_;
 
        vector<double> splitting_val_;
 
        vector<index_t> node_m_;
 
        vector<index_t> node_right_child_;
    };
 
    /*********************************************************************/
}
 
#endif