BallTree.hpp

#pragma once
 
#include <queue>
#include <memory>
#include <utility>
#include <Eigen/Core>
#include "Features.hpp"
#include "dll.hpp"
 
namespace ml
{
    class BallTree
    {
    public:
        DLL_DECLSPEC BallTree(Eigen::Ref<const Eigen::MatrixXd> X, unsigned int min_split_size);
 
        DLL_DECLSPEC BallTree(Eigen::Ref<const Eigen::MatrixXd> X, Eigen::Ref<const Eigen::VectorXd> y, unsigned int min_split_size);
 
        DLL_DECLSPEC BallTree(Eigen::MatrixXd&& X, unsigned int min_split_size);
 
        DLL_DECLSPEC BallTree(Eigen::MatrixXd&& X, Eigen::VectorXd&& y, unsigned int min_split_size);
 
        const Eigen::MatrixXd& data() const
        {
            return data_;
        }        
 
        const Eigen::VectorXd& labels() const
        {
            return labels_;
        }
 
        DLL_DECLSPEC void find_k_nearest_neighbours(Eigen::Ref<const Eigen::VectorXd> x, unsigned int k, std::vector<unsigned int>& nn) const;        
 
        DLL_DECLSPEC unsigned int find_nearest_neighbour(Eigen::Ref<const Eigen::VectorXd> x) const;
 
        auto size() const
        {
            return static_cast<unsigned int>(data_.cols());
        }
 
        auto dim() const
        {
            return static_cast<unsigned int>(data_.rows());
        }
    private:
        struct Node
        {
            double radius; 
            unsigned int pivot_index; 
            unsigned int start_index; 
            unsigned int end_index; 
            std::unique_ptr<Node> left_child;
            std::unique_ptr<Node> right_child;
        };
 
        Eigen::MatrixXd data_; 
        Eigen::VectorXd labels_; 
        std::unique_ptr<Node> root_; 
        unsigned int min_split_size_;
 
        void construct(Eigen::Ref<Eigen::MatrixXd> work, Eigen::Ref<Eigen::VectorXd> labels, unsigned int offset, std::unique_ptr<Node>& node, Features::VectorRange<Features::IndexedFeatureValue> features);
 
        typedef std::pair<unsigned int, double> IndexedDistanceFromTarget;
 
        struct IndexedDistanceFromTargetComparator
        {
            bool operator()(const IndexedDistanceFromTarget& a, const IndexedDistanceFromTarget& b) const
            {
                return a.second < b.second;
            }
        };
 
        typedef std::priority_queue<IndexedDistanceFromTarget, std::vector<IndexedDistanceFromTarget>, IndexedDistanceFromTargetComparator> MaxDistancePriorityQueue;
 
        void knn_search(Eigen::Ref<const Eigen::VectorXd> x, unsigned int k, const Node* node, MaxDistancePriorityQueue& q) const;
 
        double distance_from_queue(Eigen::Ref<const Eigen::VectorXd> x, unsigned int k, const MaxDistancePriorityQueue& q) const;
    };
}