KD_Tree.h

Go to the documentation of this file.

#ifndef   dsa_KD_Tree_H__
#define   dsa_KD_Tree_H__

#include "../utility.h"
#include "../math/utility.h"

#include <cstdlib>

#include <vector>
#include <algorithm>
#include <queue>

namespace meow {

template<class Vector, class Scalar>
class KD_Tree {
private:
  struct Node {
    Vector  vector_;
    ssize_t lChild_;
    ssize_t rChild_;

    Node(Vector v, ssize_t l, ssize_t r): vector_(v), lChild_(l), rChild_(r){
    }
  };
  typedef std::vector<Node> Nodes;

  class Sorter {
  private:
    Nodes const* nodes_;
    size_t       cmp_;
  public:
    Sorter(Nodes const* nodes, size_t cmp):
    nodes_(nodes), cmp_(cmp){
    }
    bool operator()(size_t const& a, size_t const& b) const{
      if ((*nodes_)[a].vector_[cmp_] != (*nodes_)[b].vector_[cmp_]) {
        return ((*nodes_)[a].vector_[cmp_] < (*nodes_)[b].vector_[cmp_]);
      }
      return ((*nodes_)[a].vector_ < (*nodes_)[b].vector_);
    }
  };
  struct Answer {
    ssize_t index_;
    Scalar  dist2_;
    //
    Answer(ssize_t index, Scalar dist2):
    index_(index), dist2_(dist2) {
    }
    Answer(Answer const& answer2):
    index_(answer2.index_), dist2_(answer2.dist2_) {
    }
  };
  class AnswerCompare {
  private:
    Nodes const* nodes_;
    bool         cmpValue_;
  public:
    AnswerCompare(Nodes const* nodes, bool cmpValue):
    nodes_(nodes), cmpValue_(cmpValue) {
    }
    bool operator()(Answer const& a, Answer const& b) const {
      if (cmpValue_ == true && a.dist2_ == b.dist2_) {
        return ((*nodes_)[a.index_].vector_ < (*nodes_)[b.index_].vector_);
      }
      return (a.dist2_ < b.dist2_);
    }
  };
  typedef std::vector<Answer> AnswerV;
  typedef std::priority_queue<Answer, AnswerV, AnswerCompare> Answers;
  //
  const ssize_t kNIL_;
  //
  Nodes        nodes_;
  size_t        root_;
  bool   needRebuild_;
  size_t   dimension_;
  //
  Scalar distance2(Vector const& v1, Vector const& v2) const {
    Scalar ret(0);
    for(size_t i = 0; i < dimension_; i++){
      ret += squ(v1[i] - v2[i]);
    }
    return ret;
  }
  //
  void query(Vector const&        v,
             size_t               nearestNumber,
             AnswerCompare const& answerCompare,
             ssize_t              index,
             int                  depth,
             std::vector<Scalar>& dist2Vector,
             Scalar               dist2Minimum,
             Answers             *out) const {
    if (index == kNIL_) return ;
    size_t cmp = depth % dimension_;
    ssize_t this_side, that_side;
    if (!(nodes_[index].vector_[cmp] < v[cmp])) {
      this_side = nodes_[index].lChild_;
      that_side = nodes_[index].rChild_;
    }else{
      this_side = nodes_[index].rChild_;
      that_side = nodes_[index].lChild_;
    }
    query(v, nearestNumber, answerCompare,
          this_side, depth + 1,
          dist2Vector, dist2Minimum,
          out);
    Answer my_ans(index, distance2(nodes_[index].vector_, v));
    if (out->size() < nearestNumber || answerCompare(my_ans, out->top())) {
      out->push(my_ans);
      if (out->size() > nearestNumber) out->pop();
    }
    Scalar dist2_old(dist2Vector[cmp]);
    dist2Vector[cmp] = squ(nodes_[index].vector_[cmp] - v[cmp]);
    Scalar dist2Minimum2(dist2Minimum + dist2Vector[cmp] - dist2_old);
    if (out->size() < nearestNumber || !(out->top().dist2_ < dist2Minimum)) {
      query(v, nearestNumber, answerCompare,
            that_side, depth + 1,
            dist2Vector, dist2Minimum2,
            out);
    }
    dist2Vector[cmp] = dist2_old;
  }
  ssize_t build(ssize_t              beg,
                ssize_t              end,
                std::vector<size_t>* orders,
                int                  depth) {
    if (beg > end) return kNIL_;
    size_t tmp_order  = dimension_;
    size_t which_side = dimension_ + 1;
    ssize_t mid = (beg + end) / 2;
    size_t  cmp = depth % dimension_;
    for (ssize_t i = beg; i <= mid; i++) {
      orders[which_side][orders[cmp][i]] = 0;
    }
    for (ssize_t i = mid + 1; i <= end; i++) {
      orders[which_side][orders[cmp][i]] = 1;
    }
    for (size_t i = 0; i < dimension_; i++) {
      if (i == cmp) continue;
      size_t left = beg, right = mid + 1;
      for (int j = beg; j <= end; j++) {
        size_t ask = orders[i][j];
        if(ask == orders[cmp][mid]) {
          orders[tmp_order][mid] = ask;
        }
        else if(orders[which_side][ask] == 1) {
          orders[tmp_order][right++] = ask;
        }
        else {
          orders[tmp_order][left++] = ask;
        }
      }
      for (int j = beg; j <= end; j++) {
        orders[i][j] = orders[tmp_order][j];
      }
    }
    nodes_[orders[cmp][mid]].lChild_ = build(beg, mid - 1, orders, depth + 1);
    nodes_[orders[cmp][mid]].rChild_ = build(mid + 1, end, orders, depth + 1);
    return orders[cmp][mid];
  }
public:
  typedef typename std::vector<Vector> Vectors;

  KD_Tree(): kNIL_(-1), root_(kNIL_), needRebuild_(false), dimension_(1) {
  }

  KD_Tree(size_t dimension):
  kNIL_(-1), root_(kNIL_), needRebuild_(false), dimension_(dimension) {
  }

  ~KD_Tree() {
  }

  void insert(Vector const& v) {
    nodes_.push_back(Node(v, kNIL_, kNIL_));
    needRebuild_ = true;
  }

  bool erase(Vector const& v) {
    for (size_t i = 0, I = nodes_.size(); i < I; i++) {
      if (nodes_[i] == v) {
        if (i != I - 1) {
          std::swap(nodes_[i], nodes_[I - 1]);
        }
        needRebuild_ = true;
        return true;
      }
    }
    return false;
  }

  void build(){
    if (needRebuild_) {
      forceBuild();
    }
  }

  void forceBuild() {
    std::vector<size_t> *orders = new std::vector<size_t>[dimension_ + 2];
    for (size_t j = 0; j < dimension_ + 2; j++) {
      orders[j].resize(nodes_.size());
    }
    for (size_t j = 0; j < dimension_; j++) {
      for (size_t i = 0, I = nodes_.size(); i < I; i++) {
        orders[j][i] = i;
      }
      std::sort(orders[j].begin(), orders[j].end(), Sorter(&nodes_, j));
    }
    root_ = build(0, (ssize_t)nodes_.size() - 1, orders, 0);
    delete [] orders;
    needRebuild_ = false;
  }

  Vectors query(Vector const& v,
                size_t        nearestNumber,
                bool          compareWholeVector) const {
    ((KD_Tree*)this)->build();
    AnswerCompare answer_compare(&nodes_, compareWholeVector);
    Answers       answer_set(answer_compare);
    std::vector<Scalar> tmp(dimension_, 0);
    query(v, nearestNumber,
          answer_compare,
          root_, 0,
          tmp, Scalar(0),
          &answer_set);
    Vectors ret(answer_set.size());
    for (int i = (ssize_t)answer_set.size() - 1; i >= 0; i--) {
      ret[i] = nodes_[answer_set.top().index_].vector_;
      answer_set.pop();
    }
    return ret;
  }

  void clear() {
    root_ = kNIL_;
    nodes_.clear();
    needRebuild_ = false;
  }

  void reset(size_t dimension) {
    clear();
    dimension_ = dimension;
  }
};

} // meow

#endif // dsa_KD_Tree_H__