VP_Tree.h

Go to the documentation of this file.

#ifndef   dsa_VP_Tree_H__
#define   dsa_VP_Tree_H__

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

#include <cstdlib>

#include <list>
#include <vector>
#include <stack>
#include <queue>

namespace meow {

template<class Vector, class Scalar>
class VP_Tree {
public:
  typedef std::vector<Vector> Vectors;
private:
  struct Node {
    size_t     index_;
    Scalar threshold_;
    Node*  nearChild_;
    Node*   farChild_;
    //
    Node(size_t index): index_(index), nearChild_(NULL), farChild_(NULL){
    }
  };
  struct Answer {
    size_t index_;
    Scalar dist2_;
    //
    Answer(size_t index, Scalar const& dist2): index_(index), dist2_(dist2){
    }
    Answer(Answer const& answer2):
    index_(answer2.index_), dist2_(answer2.dist2_){
    }
  };
  class AnswerCompare {
  private:
    Vectors const*  vectors_;
    bool           cmpValue_;
  public:
    AnswerCompare(Vectors const* vectors, bool cmpValue):
    vectors_(vectors), cmpValue_(cmpValue){
    }
    bool operator()(Answer const& a, Answer const& b) const {
      if (a.dist2_ < b.dist2_) return true;
      if (b.dist2_ < a.dist2_) return false;
      return (cmpValue_ && ((*vectors_)[a.index_] < (*vectors_)[b.index_]));
    }
  };
  typedef std::vector<Answer> AnswerV;
  typedef std::priority_queue<Answer, AnswerV, AnswerCompare> Answers;

  Vectors     vectors_;
  Node*          root_;
  size_t    dimension_;
  bool    needRebuild_;

  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;
  }
  int distanceCompare(Scalar const& a2, Scalar const& b2,
                      Scalar const& c2) const {
    if (b2 < 0) {
      return -distanceCompare(c2, -b2, a2);
    }
    Scalar cab(c2 - a2 - b2);
    if (cab < Scalar(0)) return 1;
    Scalar ab2(Scalar(4) * a2 * b2), cab2(squ(cab));
    if      ( ab2 < cab2) return -1;
    else if (cab2 <  ab2) return  1;
    else                  return  0;
  }
  Scalar split(ssize_t first, ssize_t last, size_t order,
               Vector const& center) {
    ssize_t first0 = first;
    std::vector<Scalar> dist2(last - first + 1);
    for (ssize_t i = first; i <= last; i++) {
      dist2[i - first0] = distance2(vectors_[i], center);
    }
    while (first < last) {
      size_t thresholdindex_ = first + rand() % (last - first + 1);
      Scalar threshold(dist2[thresholdindex_ - first0]);
      size_t large_first = last + 1;
      for( ssize_t i=first; first<=(ssize_t)large_first-1; large_first--) {
        if (threshold < dist2[large_first - 1 - first0]) continue;
        while (i < (ssize_t)large_first-1&&!(threshold < dist2[i-first0])) i++;
        if (i < (ssize_t)large_first - 1){
          std::swap(dist2   [large_first - 1 - first0], dist2   [i - first0]);
          std::swap(vectors_[large_first - 1         ], vectors_[i         ]);
          i++;
        }
        else {
          break;
        }
      }
      if (large_first == (size_t)last + 1) {
        std::swap(dist2   [thresholdindex_-first0], dist2   [last-first0]);
        std::swap(vectors_[thresholdindex_       ], vectors_[last       ]);
        if ((ssize_t)order == last - first) {
          first = last;
          break;
        }
        last--;
      }
      else {
        if (order < large_first - first) {
          last = large_first - 1;
        }
        else {
          order -= large_first - first;
          first = large_first;
        }
      }
    }
    return dist2[first - first0];
  }
  //
  Node* build(ssize_t first, ssize_t last) {
    if (first > last) return NULL;
    Node* ret = new Node(first);
    if (first < last) {
      std::swap(vectors_[first],
                vectors_[first + rand() % (last - first + 1)]);
      ssize_t mid = (first + 1 + last + 1) / 2;
      ret->threshold_ = split(first + 1, last, mid - (first + 1),
                              vectors_[first]);
      ret->nearChild_ = build(first + 1, mid - 1      );
      ret->farChild_  = build(           mid    , last);
    }
    return ret;
  }
  void query(Vector const&        vector,
             size_t               k,
             AnswerCompare const& cmp,
             Node const*          node,
             Answers*             out) const {
    if (node == NULL) return ;
    Scalar dist2 = distance2(vector, vectors_[node->index_]);
    Answer my_ans(node->index_, dist2);
    if (out->size() < k || cmp(my_ans, out->top())) {
      out->push(my_ans);
      if (out->size() > k) {
        out->pop();
      }
    }
    if (node->nearChild_ == NULL && node->farChild_ == NULL) return ;
    if (out->size() < k || distanceCompare(dist2, -out->top().dist2_,
                                              node->threshold_) <= 0) {
      query(vector, k, cmp, node->nearChild_, out);
    }
    if (out->size() < k || distanceCompare(dist2,  out->top().dist2_,
                                              node->threshold_) >= 0) {
      query(vector, k, cmp, node->farChild_, out);
    }
  }
  void clear(Node* root) {
    if(root == NULL) return ;
    clear(root->nearChild_);
    clear(root->farChild_);
    delete root;
  }
  Node* dup(Node* root) {
    if(root == NULL) return ;
    Node* ret = new Node(root->index_);
    ret->threshold_ = root->threshold_;
    ret->nearChild_ = dup(root->nearChild_);
    ret->farChild_  = dup(root->farChild_ );
    return ret;
  }
public:
  VP_Tree(): root_(NULL), vectors_(0), dimension_(1), needRebuild_(false){
    reset(0);
  }

  VP_Tree(VP_Tree const& tree2):
  vectors_(tree2.vectors_),
  root_(dup(tree2.root_)),
  dimension_(tree2.dimension_),
  needRebuild_(tree2.needRebuild_) {
  }

  VP_Tree(size_t dimension):
  vectors_(0),
  root_(NULL),
  dimension_(0),
  needRebuild_(false) {
    reset(dimension);
  }

  ~VP_Tree() {
    clear(root_);
  }

  VP_Tree& copyFrom(VP_Tree const& tree2) {
    reset(tree2.dimension_);
    vectors_     = tree2.vectors_;
    root_        = dup(tree2.root_);
    needRebuild_ = tree2.needRebuild_;
    return *this;
  }

  void insert(Vector const& vector) {
    vectors_.push_back(vector);
    needRebuild_ = true;
  }

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

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

  void forceBuild() {
    root_ = build(0, (size_t)vectors_.size() - 1);
    needRebuild_ = false;
  }

  Vectors query(Vector const& vector,
                size_t        nearestNumber,
                bool          compareWholeVector) const {
    ((VP_Tree*)this)->build();
    AnswerCompare cmp(&vectors_, compareWholeVector);
    Answers answers(cmp);
    query(vector, nearestNumber, cmp, root_, &answers);
    std::stack<Answer> rev;
    for ( ; !answers.empty(); answers.pop()) rev.push(answers.top());
    Vectors ret;
    for ( ; !rev.empty(); rev.pop()) ret.push_back(vectors_[rev.top().index_]);
    return ret;
  }

  void clear() {
    clear(root_);
    vectors_.clear();
    root_ = NULL;
    needRebuild_ = false;
  }

  size_t reset(size_t dimension) {
    clear();
    dimension_ = std::max((size_t)1, dimension);
    return dimension_;
  }

  VP_Tree& operator=(VP_Tree const& tree2) {
    return copyFrom(tree2);
  }
};

} // meow

#endif // dsa_VP_Tree_H__