c++实现kd树

#ifndef _KD_TREE_H_
#define _KD_TREE_H_

#include <memory>
#include <vector>
#include <algorithm>
#include <iostream>
#include <functional>
#include <iomanip>
#include <stack>
#include <array>
#include <cfloat>
#include <cmath>

namespace zstd
{
    struct threeD_node
    {
        double value[3];//value[0] = x, value[1] = y, value[2] = z
        threeD_node()
        {
            value[0] = 0.0;
            value[1] = 0.0;
            value[2] = 0.0;
        }
        threeD_node(double x, double y, double z)
        {
            value[0] = x;
            value[1] = y;
            value[2] = z;
        }
    };
    struct sort_for_threeD_node
    {
        int dimension;
        sort_for_threeD_node(int d) :dimension(d){}
        bool operator ()(const threeD_node& lhs, const threeD_node& rhs)
        {
            if (dimension == 0)
                return lhs.value[0] < rhs.value[0];
            else if (dimension == 1)
                return lhs.value[1] < rhs.value[1];
            else if (dimension == 2)
                return lhs.value[2] < rhs.value[2];
            else
                std::cerr << "error in sort_for_threeD_node"<< std::endl;
            return false;
        }
    };

    struct kd_node
    {
        double value[3];//value[0] = x, value[1] = y, value[2] = z
        int kv;//0=x, 1=y, 2=z
        bool is_leaf;
        kd_node  *left, *right;
        kd_node()
        {
            value[0] = value[1] = value[2] = 0.0;
            kv = -1;
            is_leaf = false;
            left = nullptr;
            right = nullptr;
        }
        kd_node(const kd_node& node)
        {
            value[0] = node.value[0];
            value[1] = node.value[1];
            value[2] = node.value[2];
            kv = node.kv;
            is_leaf = node.is_leaf;
            left = node.left;
            right = node.right;
        }
        kd_node& operator = (const kd_node& node)
        {
            value[0] = node.value[0];
            value[1] = node.value[1];
            value[2] = node.value[2];
            kv = node.kv;
            is_leaf = node.is_leaf;
            left = node.left;
            right = node.right;

            return *this;
        }
    };
    class kd_tree
    {
    private:
        std::shared_ptr<kd_node> root;
        std::vector<threeD_node>& vec_ref;
        const int k = 3;
        const int cspace = 4;
    private:
        int get_dimension(int n) const
        {
            return n % k;
        }
        void sort_by_dimension(std::vector<threeD_node>& v, int dimension, int l, int r);
        kd_node* build_tree(int left, int right, kd_node* sp_node, int dimension);
        void _print_tree(kd_node* sp, bool left, int space);

        double distance(const kd_node& lhs, const threeD_node& rhs);
    public:
        explicit kd_tree(std::vector<threeD_node>&);
        kd_tree(const kd_tree&) = delete;
        kd_tree operator = (const kd_tree&) = delete;
        ~kd_tree(){};

        void print_tree();
        std::vector<threeD_node> find_k_nearest(int k, const threeD_node& D);
    };
    void kd_tree::sort_by_dimension(std::vector<threeD_node>& v, int dimension, int l, int r)
    {
        sort_for_threeD_node s(dimension);
        std::sort(v.begin()+l, v.begin()+r, s);
    }
    kd_tree::kd_tree(std::vector<threeD_node>& v) :vec_ref(v)
    {
        if (vec_ref.empty())
            root = nullptr;
        else
        {
            root = std::make_shared<kd_node>();
            int dimension = 0;
            sort_by_dimension(vec_ref, dimension, 0, vec_ref.size());
            int mid = vec_ref.size() / 2;
            root->value[0] = vec_ref[mid].value[0];
            root->value[1] = vec_ref[mid].value[1];
            root->value[2] = vec_ref[mid].value[2];
            root->kv = dimension;
            if (vec_ref.size() == 1)//root is leaf
            {
                root->left = nullptr;
                root->right = nullptr;
                root->is_leaf = true;
            }
            else
            {
                root->is_leaf = false;
                root->left = build_tree(0, mid - 1, root->left, get_dimension(dimension + 1));
                root->right = build_tree(mid + 1, vec_ref.size() - 1, root->right, get_dimension(dimension + 1));
            }
        }
    }
    kd_node* kd_tree::build_tree(int left, int right, kd_node* sp_node, int dimension)
    {
        dimension = get_dimension(dimension);
        sort_by_dimension(vec_ref, dimension, left, right + 1);

        if(left == right)//leaf
        {
            sp_node = new kd_node();
            sp_node->value[0] = vec_ref[left].value[0];
            sp_node->value[1] = vec_ref[left].value[1];
            sp_node->value[2] = vec_ref[left].value[2];
            sp_node->kv = dimension;
            sp_node->is_leaf = true;
            sp_node->left = nullptr;
            sp_node->right = nullptr;

            return sp_node;
        }
        else if (left < right)
        {
            int mid = left + (right - left) / 2;
            sp_node = new kd_node();
            sp_node->value[0] = vec_ref[mid].value[0];
            sp_node->value[1] = vec_ref[mid].value[1];
            sp_node->value[2] = vec_ref[mid].value[2];
            sp_node->kv = dimension;
            sp_node->is_leaf = false;
            sp_node->left = nullptr;
            sp_node->right = nullptr;

            sp_node->left = build_tree(left, mid - 1, sp_node->left, get_dimension(dimension + 1));
            sp_node->right = build_tree(mid + 1, right, sp_node->right, get_dimension(dimension + 1));

            return sp_node;
        }
        return nullptr;
    }
    void kd_tree::_print_tree(kd_node* sp, bool left, int space)
    {
        if (sp != nullptr)
        {
            _print_tree(sp->right, false, space + cspace);
            std::cout << std::setw(space);
            std::cout << "(" <<
                sp->value[0] << ", " <<
                sp->value[1] << ", " <<
                sp->value[2] << ")";
            if (left)
                std::cout << "left";
            else
                std::cout << "right";
            if (sp->is_leaf)
                std::cout << "------leaf";
            std::cout << std::endl;
            _print_tree(sp->left, true, space + cspace);
        }
        else
            std::cout << std::endl;
    }
    void kd_tree::print_tree()
    {
        std::cout << "kd_tree : " << std::endl;
        if (root != nullptr)
        {
            int space = 0;
            _print_tree(root->right, false, space + cspace);
            std::cout << "(" << 
                root->value[0] << ", " << 
                root->value[1] << ", " << 
                root->value[2] << ")root" << std::endl;
            _print_tree(root->left, true, space + cspace);
        }
    }
    double kd_tree::distance(const kd_node& lhs, const threeD_node& rhs)
    {
        double v0 = lhs.value[0] - rhs.value[0];
        double v1 = lhs.value[1] - rhs.value[1];
        double v2 = lhs.value[2] - rhs.value[2];
        return sqrt(v0 * v0 + v1 * v1 + v2 * v2);
    }
    std::vector<threeD_node> kd_tree::find_k_nearest(int ks, const threeD_node& D)
    {
        std::vector<threeD_node> res;
        const kd_node *ptr_kd_node;
        if (static_cast<std::size_t>(ks) > vec_ref.size())
            return res;
        std::stack<kd_node> s;
        struct pair
        {
            double distance;
            kd_node node;
            pair() :distance(DBL_MAX), node(){ }
            bool operator < (const pair& rhs)
            {
                return distance < rhs.distance;
            }
        };
        std::unique_ptr<pair[]> ptr_pair(new pair[ks]);
        //pair *ptr_pair = new pair[ks]();
        if (!ptr_pair)
            exit(-1);

        if (!root)//the tree is empty
            return std::vector<threeD_node>();
        else
        {
            if (D.value[root->kv] < root->value[root->kv])
            {
                s.push(*root);
                ptr_kd_node = root->left;
            }
            else
            {
                s.push(*root);
                ptr_kd_node = root->right;
            }
            while (ptr_kd_node != nullptr)
            {
                if (D.value[ptr_kd_node->kv] < ptr_kd_node->value[ptr_kd_node->kv])
                {
                    s.push(*ptr_kd_node);
                    ptr_kd_node = ptr_kd_node->left;
                }
                else
                {
                    s.push(*ptr_kd_node);
                    ptr_kd_node = ptr_kd_node->right;
                }
            }
            
            while (!s.empty())
            {
                kd_node popped_kd_node;//±£´æ×îеĴÓÕ»ÖÐpop³öµÄkd_node
                popped_kd_node = s.top();
                s.pop();
                double dist = distance(popped_kd_node, D);
                std::sort(&ptr_pair[0], &ptr_pair[ks]);
                if (dist < ptr_pair[ks-1].distance)
                {
                    ptr_pair[ks-1].distance = dist;
                    ptr_pair[ks-1].node = popped_kd_node;
                }

                if (abs(D.value[popped_kd_node.kv] - popped_kd_node.value[popped_kd_node.kv])
                        >= dist)//Ô²²»ºÍpopped_kd_nodeµÄÁíÒ»°ëÇøÓòÏཻ
                    continue;
                else//Ô²ºÍpopped_kd_nodeµÄÁíÒ»°ëÇøÓòÏཻ
                {
                    if (D.value[popped_kd_node.kv] < popped_kd_node.value[popped_kd_node.kv])//right
                    {
                        kd_node *ptr = popped_kd_node.right;
                        while (ptr != nullptr)
                        {    
                            s.push(*ptr);
                            if (D.value[ptr->kv] < ptr->value[ptr->kv])
                                ptr = ptr->left;
                            else
                                ptr = ptr->right;
                        }
                    }
                    else//left
                    {
                        kd_node *ptr = popped_kd_node.left;
                        while (ptr != nullptr)
                        {
                            s.push(*ptr);
                            if (D.value[ptr->kv] < ptr->value[ptr->kv])
                                ptr = ptr->left;
                            else
                                ptr = ptr->right;
                        }
                    }
                }
            }//end of while
            for(int i = 0; i != ks; ++i)
                res.push_back(threeD_node(ptr_pair[i].node.value[0], 
                            ptr_pair[i].node.value[1], ptr_pair[i].node.value[2]));
        }//end of else
        //delete ptr_pair;
        return res;
    }

}//end of namespace zstd

#endif

#include <string>
#include <iostream>
#include <new>
#include <fstream>
#include <vector>
#include <algorithm>
#include <ctime>

#include "trie_tree.h"
#include "kd_tree.h"

int main()
{
    std::vector<zstd::threeD_node> v, res;
    v.push_back(zstd::threeD_node(2, 3, 1));//14
    v.push_back(zstd::threeD_node(5, 4, 7));//90
    v.push_back(zstd::threeD_node(9, 6, 9));//198
    v.push_back(zstd::threeD_node(4, 7, 2));//69
    v.push_back(zstd::threeD_node(8, 1, 5));//90
    v.push_back(zstd::threeD_node(7, 2, 0));//53
    v.push_back(zstd::threeD_node(8, 8, 8));//192
    v.push_back(zstd::threeD_node(1, 2, 3));//14
    v.push_back(zstd::threeD_node(5, 2, 1));//30
    v.push_back(zstd::threeD_node(12, 23, 0));//673
    v.push_back(zstd::threeD_node(10, 0, 2));//104
    std::cout << "size: " << v.size() << std::endl;
    zstd::kd_tree tree(v);
    tree.print_tree();
    res = tree.find_k_nearest(11, zstd::threeD_node(0, 0, 0));
    std::cout << "-------" << std::endl;
    std::cout << "离点(0,0,0)最近的点依次是：" << std::endl;
    for (auto i : res)
    {
        std::cout << "(" << i.value[0] << ", " << i.value[1] << ", " << i.value[2] << ")" << std::endl;
    }
    system("pause");
    return 0;
}