KDTree.h

Go to the documentation of this file.

/*
 * $Id: KDTree.h 5871 2005-10-28 02:10:33Z dfs $
 *
 * Copyright 2003-2005 Daniel F. Savarese
 * Copyright 2005 Savarese Software Research
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.savarese.com/software/ApacheLicense-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef __SAVA_SPATIAL_KDTREE_H
#define __SAVA_SPATIAL_KDTREE_H

#include <algorithm>
#include <utility>
#include <vector>

#include <libsava/spatial/detail/KDTree.h>
#include <libsava/spatial/Point.h>

__BEGIN_PACKAGE_SAVA_SPATIAL

template<typename Tree>
struct KDTreeTraits {
  typedef typename Tree::key_type key_type;
  typedef typename Tree::mapped_type mapped_type;
  typedef typename Tree::value_type value_type;
  typedef typename Tree::pointer pointer;
  typedef typename Tree::const_pointer const_pointer;
  typedef typename Tree::reference reference;
  typedef typename Tree::const_reference const_reference;
  typedef typename Tree::discriminator_type discriminator_type;
  typedef typename Tree::node_type node_type;
  typedef typename node_type::child_type child_type;
  typedef typename Tree::iterator iterator;
  typedef typename Tree::const_iterator const_iterator;
  typedef typename Tree::point_traits point_traits;
  typedef Tree tree_type;

  static inline const unsigned int dimensions() {
    return Tree::dimensions();
  }
};


template<typename Tree>
struct KDTreeConstTraits : public KDTreeTraits<Tree> {
  typedef typename Tree::const_pointer pointer;
  typedef typename Tree::const_reference reference;
};


// Note: we store the discriminator in each node to avoid modulo division,
// trading space for time.
template <typename Point,
          typename Value,
          unsigned int Dimensions = 2,
          typename Discriminator = unsigned char,
          typename Size = unsigned int>
class KDTree {

public:

  typedef KDTreeTraits<KDTree> traits;
  typedef KDTreeConstTraits<KDTree> const_traits;
  typedef Point key_type;
  typedef Value mapped_type;
  typedef std::pair<const key_type, mapped_type> value_type;
  typedef value_type* pointer;
  typedef pointer const const_pointer;
  typedef value_type& reference;
  typedef const reference const_reference;
  typedef Discriminator discriminator_type;
  typedef detail::KDTreeRangeSearchIterator<traits> iterator;
  typedef detail::KDTreeRangeSearchIterator<const_traits> const_iterator;
  typedef PointTraits<key_type> point_traits;

  typedef Size size_type;

  static inline const unsigned int dimensions() {
    return Dimensions;
  }

  typedef detail::KDTreeNode<traits> node_type;
  typedef typename node_type::child_type child_type;

private:

  struct NodeComparator {
    discriminator_type discriminator;

    explicit NodeComparator() : discriminator(0) { }

    bool operator()(const node_type* const & n1,
                    const node_type* const & n2) const
    {
      return (n1->point()[discriminator] < n2->point()[discriminator]);
    }
  };

  node_type *_root;
  size_type _size;
  iterator _endIterator;
  const_iterator _constEndIterator;

  node_type * getNode(const key_type & point, node_type **parent = 0) const {
    discriminator_type discriminator;
    child_type child;
    node_type *node = _root, *last = 0;

    while(node != 0) {
      discriminator = node->discriminator();

      if(point[discriminator] > node->point()[discriminator])
        child = node_type::ChildHigh;
      else if(point[discriminator] < node->point()[discriminator])
        child = node_type::ChildLow;
      else if(node->point() == point) {
        if(parent != 0)
          *parent = last;
        return node;
      } else
        child = node_type::ChildHigh;

      last = node;
      node = node->child(child);
    }

    if(parent != 0)
      *parent = last;

    return 0;
  }

  node_type * getMinimumNode(node_type *node, node_type *p,
                             const discriminator_type discriminator,
                             node_type **parent)
  {
    node_type *result;

    if(discriminator == node->discriminator()) {
      if(node->child(node_type::ChildLow) != 0)
        return
          getMinimumNode(node->child(node_type::ChildLow), node,
                              discriminator, parent);
      else
        result = node;
    } else {
      node_type *nlow = 0, *nhigh = 0;
      node_type *plow, *phigh;

      if(node->child(node_type::ChildLow) != 0)
        nlow =
          getMinimumNode(node->child(node_type::ChildLow), node,
                              discriminator, &plow);

      if(node->child(node_type::ChildHigh) != 0)
        nhigh =
          getMinimumNode(node->child(node_type::ChildHigh), node,
                              discriminator, &phigh);

      if(nlow != 0 && nhigh != 0) {
        if(nlow->point()[discriminator] < nhigh->point()[discriminator]) {
          result  = nlow;
          *parent = plow;
        } else {
          result  = nhigh;
          *parent = phigh;
        }
      } else if(nlow != 0) {
        result  = nlow;
        *parent = plow;
      } else if(nhigh != 0) {
        result  = nhigh;
        *parent = phigh;
      } else
        result  = node;
    }

    if(result == node)
      *parent = p;
    else if(node->point()[discriminator] < result->point()[discriminator]) {
      result  = node;
      *parent = p;
    }

    return result;
  }


  node_type * recursiveRemoveNode(node_type *node) {
    discriminator_type discriminator;
    node_type *newRoot, *parent;

    if(node->child(node_type::ChildLow) == 0 &&
       node->child(node_type::ChildHigh) == 0)
      return 0;
    else
      discriminator = node->discriminator();

    if(node->child(node_type::ChildHigh) == 0) {
      node->child(node_type::ChildHigh) = node->child(node_type::ChildLow);
      node->child(node_type::ChildLow)  = 0;
    }

    newRoot =
      getMinimumNode(node->child(node_type::ChildHigh), node,
                     discriminator, &parent);

    child_type child = (parent->child(node_type::ChildLow) == newRoot ?
                        node_type::ChildLow : node_type::ChildHigh);
    parent->child(child) = recursiveRemoveNode(newRoot);

    newRoot->child(node_type::ChildLow)  = node->child(node_type::ChildLow);
    newRoot->child(node_type::ChildHigh) = node->child(node_type::ChildHigh);
    newRoot->discriminator() = node->discriminator();

    return newRoot;
  }


  bool add(const key_type & point, const mapped_type & value,
            node_type **node, node_type *parent,
            mapped_type *replaced = 0)
  {
    if(parent == 0) {
      if(_root != 0)
        *node = _root;
      else {
        _root = *node = new node_type(0, point, value);
        ++_size;
        return false;
      }
    } else if(*node == 0) {
      discriminator_type discriminator;
      child_type child;

      discriminator = parent->discriminator();
      child = (point[discriminator] >= parent->point()[discriminator] ?
               node_type::ChildHigh : node_type::ChildLow);

      if(++discriminator >= dimensions())
        discriminator = 0;

      parent->child(child) = *node =
        new node_type(discriminator, point, value);

      ++_size;
      return false;
    }

    if(replaced != 0)
      *replaced = (*node)->value();

    (*node)->value() = value;

    return true;
  }

  template<template<typename> class Container>
  static inline
  node_type * optimize(typename Container<node_type*>::iterator begin,
                       typename Container<node_type*>::iterator end,
                       NodeComparator & comparator)
  {
    node_type *midpoint = 0;
    typename Container<node_type*>::iterator::difference_type diff;

    diff = end - begin;

    if(diff > 1) {
      discriminator_type discriminator = comparator.discriminator;
      typename Container<node_type*>::iterator nth = begin + (diff >> 1);
      typename Container<node_type*>::iterator nthprev = nth - 1;

      //nth_element(begin, nth, end, comparator);
      stable_sort(begin, end, comparator);

      // Ties go in the right subtree.
      while(nth > begin &&
            (*nth)->point()[discriminator] == 
            (*nthprev)->point()[discriminator])
        {
          --nth;
          --nthprev;
        }

      midpoint = *nth;
      midpoint->discriminator() = discriminator;

      if(++discriminator >= dimensions())
        discriminator = 0;

      comparator.discriminator = discriminator;

      // Left subtree
      midpoint->child(node_type::ChildLow) =
        optimize<Container>(begin, nth, comparator);

      comparator.discriminator = discriminator;

      // Right subtree
      midpoint->child(node_type::ChildHigh) =
        optimize<Container>(nth + 1, end, comparator);
    } else if(diff == 1) {
      midpoint = *begin;
      midpoint->discriminator() = comparator.discriminator;
      midpoint->child(node_type::ChildLow) = 0;
      midpoint->child(node_type::ChildHigh) = 0;
    }

    return midpoint;
  }

  template<template<typename> class Container>
  static inline void fillContainer(Container<node_type*> & c, node_type *node)
  {
    if(node == 0)
      return;
    c.push_back(node);
    fillContainer(c, node->child(node_type::ChildLow));
    fillContainer(c, node->child(node_type::ChildHigh));
  }

  static inline 
  void initPoint(key_type & point,
                 const typename point_traits::coordinate_type & value)
  {
    for(unsigned int i=0; i < point_traits::dimensions(); ++i)
      point[i] = value;
  }
  
  static inline const key_type upperBound() {
    key_type bound;
    initPoint(bound, point_traits::max_coordinate());
    return bound;
  }

  static inline const key_type lowerBound() {
    key_type bound;
    initPoint(bound, point_traits::min_coordinate());
    return bound;
  }

public:

  // TODO: need assignment operator.

  explicit KDTree() : _root(0), _size(0), _endIterator(), _constEndIterator()
  { }

  KDTree(const KDTree & tree) :
    _root(0), _size(0), _endIterator(), _constEndIterator()
  {
    for(const_iterator p = tree.begin(); !p.endOfRange(); ++p)
      insert(p->first, p->second);
  }

  virtual ~KDTree() { delete _root; }

  void clear() {
    delete _root;
    _root = 0;
    _size = 0;
  }

  const size_type size() const {
    return _size;
  }

  const size_type max_size() const {
    return std::numeric_limits<size_type>::max();
  }

  bool empty() const {
    return (_root == 0);
  }

  iterator begin() {
    return iterator(lowerBound(), upperBound(), _root);
  }

  const_iterator begin() const {
    return const_iterator(lowerBound(), upperBound(), _root);
  }

  iterator & end() {
    return _endIterator;
  }

  const const_iterator & end() const {
    return _constEndIterator;
  }

  bool insert(const key_type & point, const mapped_type & value,
              mapped_type *replaced = 0)
  {
    node_type *parent;
    node_type *node = getNode(point, &parent);

    return add(point, value, &node, parent, replaced);
  }

  std::pair<iterator,bool> insert(const value_type & mapping) {
    // Ideally, we'd do this all in one step, but that will have
    // to wait until we optimize the way we handle iterators.
    bool replaced;
    mapped_type existing;
    iterator value;

    replaced = insert(mapping.first, mapping.second, &existing);
    value = find(mapping.first);

    if(replaced)
      value._node->value() = existing;

    return std::pair<iterator,bool>(value,!replaced);
  }


  mapped_type & operator[](const key_type & point) {
    node_type *parent;
    node_type *node = getNode(point, &parent);

    if(node == 0)
      add(point, mapped_type(), &node, parent);

    return node->value();
  }

  bool remove(const key_type & point, mapped_type *erased = 0) {
    node_type *parent;
    node_type *node = getNode(point, &parent);
    node_type *child;

    if(node == 0)
      return false;

    if(erased != 0)
      *erased = node->value();

    child = node;
    node  = recursiveRemoveNode(child);

    if(parent == 0)
      _root = node;
    else if(child == parent->child(node_type::ChildLow))
      parent->child(node_type::ChildLow) = node;
    else
      parent->child(node_type::ChildHigh) = node;

    // Must zero children so they are not deleted by ~node_type()
    child->child(node_type::ChildLow)  = 0; 
    child->child(node_type::ChildHigh) = 0;

    --_size;
    delete child;

    return true;
  }

  size_type erase(const key_type & point) {
    return remove(point);
  }

  void erase(iterator pos) {
    remove(pos->first);
  }

  iterator begin(const key_type & lower, const key_type & upper) {
    return iterator(lower, upper, _root);
  }

  const_iterator begin(const key_type & lower, const key_type & upper) const {
    return const_iterator(lower, upper, _root);
  }

  bool get(const key_type & point, mapped_type *value = 0) const {
    node_type *node = getNode(point);

    if(node == 0)
      return false;
    else if(value != 0)
      *value = node->value();

    return true;
  }

  iterator find(const key_type & point) {
    return iterator(point, upperBound(), _root, true);
  }

  const_iterator find(const key_type & point) const {
    return const_iterator(point, upperBound(), _root, true);
  }

  // Balances the tree.  Very expensive!
  void optimize() {
    if(empty())
      return;

    typedef std::vector<node_type*> container;
    container nodes;

    nodes.reserve(size());
    fillContainer<std::vector>(nodes, _root);

    NodeComparator comparator;
    _root =
      optimize<std::vector>(nodes.begin(), nodes.end(), comparator);
  }

};


__END_PACKAGE_SAVA_SPATIAL

#endif

---

Copyright © 2003-2005 Savarese Software Research and Daniel F. Savarese. All rights reserved.