kd树 C++实现

参考：百科kd-tree

/*
 * kdtree.h
 *
 *  Created on: Mar 3, 2017
 *      Author: wxquare
 */

#ifndef KDTREE_H_
#define KDTREE_H_

#include <vector>
#include <cmath>
#include <algorithm>
#include <iostream>
#include <stack>

template<typename T>
class KdTree {
    struct kdNode {
        std::vector<T> vec;  //data
        //split attribute,-1 means leftNode，no split attribute
        int splitAttribute;
        kdNode* lChild;
        kdNode* rChild;
        kdNode* parent;

        kdNode(std::vector<T> v = { }, int split = 0, kdNode* lch = nullptr,
                kdNode* rch = nullptr, kdNode* par = nullptr) :
                vec(v), splitAttribute(split), lChild(lch), rChild(rch), parent(par) {}
    };

private:
    kdNode *root;

public:
    KdTree() {
        root = nullptr;
    }

    KdTree(std::vector<std::vector<T>>& data) {
        root = createKdTree(data);
    }


    //matrix transpose
    std::vector<std::vector<T>> transpose(std::vector<std::vector<T>>& data) {
        int m = data.size();
        int n = data[0].size();
        std::vector<std::vector<T>> trans(n, std::vector<T>(m, 0));
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                trans[i][j] = data[j][i];
            }
        }
        return trans;
    }

    //get variance of a vector
    double getVariance(std::vector<T>& vec) {
        int n = vec.size();
        double sum = 0;
        for (int i = 0; i < n; i++) {
            sum = sum + vec[i];
        }
        double avg = sum / n;
        sum = 0; //sum of squaNN
        for (int i = 0; i < n; i++) {
            sum += pow(vec[i] - avg, 2); //#include<cmath>
        }
        return sum / n;
    }

    //According to maximum variance get split attribute.
    int getSplitAttribute(const std::vector<std::vector<T>>& data) {
        int k = data.size();
        int splitAttribute = 0;
        double maxVar = getVariance(data[0]);
        for (int i = 1; i < k; i++) {
            double temp = getVariance(data[i]);
            if (temp > maxVar) {
                splitAttribute = i;
                maxVar = temp;
            }
        }
        return splitAttribute;
    }

    //find middle value
    T getSplitValue(std::vector<T>& vec) {
        std::sort(vec.begin(), vec.end());
        return vec[vec.size() / 2];
    }

    //compute distance of two vector
    static double getDistance(std::vector<T>& v1, std::vector<T>& v2) {
        double sum = 0;
        for (size_t i = 0; i < v1.size(); i++) {
            sum += pow(v1[i] - v2[i], 2);
        }
        return sqrt(sum) / v1.size();
    }

    kdNode* createKdTree(std::vector<std::vector<T>>& data) {
        //the number of samples(data)
        if (data.empty()) return nullptr;
        int n = data.size();
        if (n == 1) {
            return new kdNode(data[0], -1); //叶子节点
        }

        //get split attribute and value
        std::vector<std::vector<T>> data_T = transpose(data);
        int splitAttribute = getSplitAttribute(data_T);
        int splitValue = getSplitValue(data_T[splitAttribute]);

        //split data according splitAttribute and splitValue
        std::vector<std::vector<T>> left;
        std::vector<std::vector<T>> right;

        int flag = 0; //the first sample's splitValue become splitnode
        kdNode *splitNode;
        for (int i = 0; i < n; i++) {
            if (flag == 0 && data[i][splitAttribute] == splitValue) {
                splitNode = new kdNode(data[i]);
                splitNode->splitAttribute = splitAttribute;
                flag = 1;
                continue;
            }
            if (data[i][splitAttribute] <= splitValue) {
                left.push_back(data[i]);
            } else {
                right.push_back(data[i]);
            }
        }

        splitNode->lChild = createKdTree(left);
        splitNode->rChild = createKdTree(right);
        return splitNode;
    }

    //search nearest neighbor
    /* 参考百度百科
     * 从root节点开始，DFS搜索直到叶子节点，同时在stack中顺序存储已经访问的节点。
       如果搜索到叶子节点，当前的叶子节点被设为最近邻节点。
       然后通过stack回溯:
       如果当前点的距离比最近邻点距离近，更新最近邻节点.
       然后检查以最近距离为半径的圆是否和父节点的超平面相交.
       如果相交，则必须到父节点的另外一侧，用同样的DFS搜索法，开始检查最近邻节点。
       如果不相交，则继续往上回溯，而父节点的另一侧子节点都被淘汰，不再考虑的范围中.
       当搜索回到root节点时，搜索完成，得到最近邻节点。
     */
    std::vector<T> searchNearestNeighbor(std::vector<T>& target,kdNode* start) {
        std::vector<T> NN;
        std::stack<kdNode*> searchPath;
        kdNode* p = start;
        while (p->splitAttribute != -1) {
            searchPath.push(p);
            int splitAttribute = p->splitAttribute;
            if (target[splitAttribute] <= p->vec[splitAttribute]) {
                p = p->lChild;
            } else {
                p = p->rChild;
            }
        }
        NN = p->vec;
        double mindis = KdTree::getDistance(target, NN);

        kdNode* cur;
        double dis;
        while (!searchPath.empty()) {
            cur = searchPath.top();
            searchPath.pop();
            dis = KdTree::getDistance(target, cur->vec);
            if (dis < mindis) {
                mindis = dis;
                NN = cur->vec;
                //判断以target为中心，以dis为半径的球是否和节点的超平面相交
                if (cur->vec[cur->splitAttribute]
                        >= target[cur->splitAttribute] - dis
                        && cur->vec[cur->splitAttribute]
                                <= target[cur->splitAttribute] + dis) {
                    std::vector<T> nn = searchNearestNeighbor(target,
                            cur->lChild);
                    if (KdTree::getDistance(target, nn)
                            < KdTree::getDistance(target, NN)) {
                        NN = nn;
                    }
                }
            }
        }
        return NN;
    }

    std::vector<T> searchNearestNeighbor(std::vector<T>& target) {
        std::vector<T> NN;
        NN = searchNearestNeighbor(target, root);
        return NN;
    }

    void print(kdNode* root) {
        std::cout << "[";
        if (root->lChild) {
            std::cout << "left:";
            print(root->lChild);
        }

        if (root) {
            std::cout << "(";
            for (size_t i = 0; i < root->vec.size(); i++) {
                std::cout << root->vec[i];
                if (i != (root->vec.size() - 1))
                    std::cout << ",";
            }
            std::cout << ")";
        }

        if (root->rChild) {
            std::cout << "right:";
            print(root->rChild);
        }
        std::cout << "]";
    }

};

#endif /* KDTREE_H_ */