kdtree.hpp

Go to the documentation of this file.

/*
 * This file is part of OpenTTD.
 * OpenTTD is free software; you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, version 2.
 * OpenTTD is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with OpenTTD. If not, see <https://www.gnu.org/licenses/old-licenses/gpl-2.0>.
 */

 
#ifndef KDTREE_HPP
#define KDTREE_HPP

template <typename T, typename TxyFunc, typename CoordT, typename DistT>
class Kdtree {
  struct node {
    T      element;  
    size_t left;     
    size_t right;    
 
    node(T element) : element(element), left(INVALID_NODE), right(INVALID_NODE) { }
  };
 
  static const size_t INVALID_NODE = SIZE_MAX;     
  static const size_t MIN_REBALANCE_THRESHOLD = 8; 
 
  std::vector<node> nodes;       
  std::vector<size_t> free_list; 
  size_t root;                   
  size_t unbalanced;             

  size_t AddNode(const T &element)
  {
    if (this->free_list.empty()) {
      this->nodes.emplace_back(element);
      return this->nodes.size() - 1;
    } else {
      size_t newidx = this->free_list.back();
      this->free_list.pop_back();
      this->nodes[newidx] = node{ element };
      return newidx;
    }
  }

  template <typename It>
  CoordT SelectSplitCoord(It begin, It end, int level)
  {
    It mid = begin + (end - begin) / 2;
    std::nth_element(begin, mid, end, [&](T a, T b) { return TxyFunc()(a, level % 2) < TxyFunc()(b, level % 2); });
    return TxyFunc()(*mid, level % 2);
  }

  template <typename It>
  size_t BuildSubtree(It begin, It end, int level)
  {
    ptrdiff_t count = end - begin;
 
    if (count == 0) {
      return INVALID_NODE;
    } else if (count == 1) {
      return this->AddNode(*begin);
    } else if (count > 1) {
      CoordT split_coord = this->SelectSplitCoord(begin, end, level);
      It split = std::partition(begin, end, [&](T v) { return TxyFunc()(v, level % 2) < split_coord; });
      size_t newidx = this->AddNode(*split);
      this->nodes[newidx].left = this->BuildSubtree(begin, split, level + 1);
      this->nodes[newidx].right = this->BuildSubtree(split + 1, end, level + 1);
      return newidx;
    } else {
      NOT_REACHED();
    }
  }

  bool Rebuild(const T *include_element, const T *exclude_element)
  {
    size_t initial_count = this->Count();
    if (initial_count < MIN_REBALANCE_THRESHOLD) return false;
 
    T root_element = this->nodes[this->root].element;
    std::vector<T> elements = this->FreeSubtree(this->root);
    elements.push_back(root_element);
 
    if (include_element != nullptr) {
      elements.push_back(*include_element);
      initial_count++;
    }
    if (exclude_element != nullptr) {
      typename std::vector<T>::iterator removed = std::remove(elements.begin(), elements.end(), *exclude_element);
      elements.erase(removed, elements.end());
      initial_count--;
    }
 
    this->Build(elements.begin(), elements.end());
    assert(initial_count == this->Count());
    return true;
  }

  void InsertRecursive(const T &element, size_t node_idx, int level)
  {
    /* Dimension index of current level */
    int dim = level % 2;
    /* Node reference */
    node &n = this->nodes[node_idx];
 
    /* Coordinate of element splitting at this node */
    CoordT nc = TxyFunc()(n.element, dim);
    /* Coordinate of the new element */
    CoordT ec = TxyFunc()(element, dim);
    /* Which side to insert on */
    size_t &next = (ec < nc) ? n.left : n.right;
 
    if (next == INVALID_NODE) {
      /* New leaf */
      size_t newidx = this->AddNode(element);
      /* Vector may have been reallocated at this point, n and next are invalid */
      node &nn = this->nodes[node_idx];
      if (ec < nc) nn.left = newidx; else nn.right = newidx;
    } else {
      this->InsertRecursive(element, next, level + 1);
    }
  }

  std::vector<T> FreeSubtree(size_t node_idx)
  {
    std::vector<T> subtree_elements;
    node &n = this->nodes[node_idx];
 
    /* We'll be appending items to the free_list, get index of our first item */
    size_t first_free = this->free_list.size();
    /* Prepare the descent with our children */
    if (n.left != INVALID_NODE) this->free_list.push_back(n.left);
    if (n.right != INVALID_NODE) this->free_list.push_back(n.right);
    n.left = n.right = INVALID_NODE;
 
    /* Recursively free the nodes being collected */
    for (size_t i = first_free; i < this->free_list.size(); i++) {
      node &fn = this->nodes[this->free_list[i]];
      subtree_elements.push_back(fn.element);
      if (fn.left != INVALID_NODE) this->free_list.push_back(fn.left);
      if (fn.right != INVALID_NODE) this->free_list.push_back(fn.right);
      fn.left = fn.right = INVALID_NODE;
    }
 
    return subtree_elements;
  }

  size_t RemoveRecursive(const T &element, size_t node_idx, int level)
  {
    /* Node reference */
    node &n = this->nodes[node_idx];
 
    if (n.element == element) {
      /* Remove this one */
      this->free_list.push_back(node_idx);
      if (n.left == INVALID_NODE && n.right == INVALID_NODE) {
        /* Simple case, leaf, new child node for parent is "none" */
        return INVALID_NODE;
      } else {
        /* Complex case, rebuild the sub-tree */
        std::vector<T> subtree_elements = this->FreeSubtree(node_idx);
        return this->BuildSubtree(subtree_elements.begin(), subtree_elements.end(), level);;
      }
    } else {
      /* Search in a sub-tree */
      /* Dimension index of current level */
      int dim = level % 2;
      /* Coordinate of element splitting at this node */
      CoordT nc = TxyFunc()(n.element, dim);
      /* Coordinate of the element being removed */
      CoordT ec = TxyFunc()(element, dim);
      /* Which side to remove from */
      size_t next = (ec < nc) ? n.left : n.right;
      assert(next != INVALID_NODE); // node must exist somewhere and must be found before a leaf is reached
      /* Descend */
      size_t new_branch = this->RemoveRecursive(element, next, level + 1);
      if (new_branch != next) {
        /* Vector may have been reallocated at this point, n and next are invalid */
        node &nn = this->nodes[node_idx];
        if (ec < nc) nn.left = new_branch; else nn.right = new_branch;
      }
      return node_idx;
    }
  }
 
 
  DistT ManhattanDistance(const T &element, CoordT x, CoordT y) const
  {
    return abs((DistT)TxyFunc()(element, 0) - (DistT)x) + abs((DistT)TxyFunc()(element, 1) - (DistT)y);
  }

  using node_distance = std::pair<T, DistT>;
  static node_distance SelectNearestNodeDistance(const node_distance &a, const node_distance &b)
  {
    if (a.second < b.second) return a;
    if (b.second < a.second) return b;
    if (a.first < b.first) return a;
    if (b.first < a.first) return b;
    NOT_REACHED(); // a.first == b.first: same element must not be inserted twice
  }

  node_distance FindNearestRecursive(CoordT xy[2], size_t node_idx, int level, DistT limit = std::numeric_limits<DistT>::max()) const
  {
    /* Dimension index of current level */
    int dim = level % 2;
    /* Node reference */
    const node &n = this->nodes[node_idx];
 
    /* Coordinate of element splitting at this node */
    CoordT c = TxyFunc()(n.element, dim);
    /* This node's distance to target */
    DistT thisdist = this->ManhattanDistance(n.element, xy[0], xy[1]);
    /* Assume this node is the best choice for now */
    node_distance best = std::make_pair(n.element, thisdist);
 
    /* Next node to visit */
    size_t next = (xy[dim] < c) ? n.left : n.right;
    if (next != INVALID_NODE) {
      /* Check if there is a better node down the tree */
      best = SelectNearestNodeDistance(best, this->FindNearestRecursive(xy, next, level + 1));
    }
 
    limit = std::min(best.second, limit);
 
    /* Check if the distance from current best is worse than distance from target to splitting line,
     * if it is we also need to check the other side of the split. */
    size_t opposite = (xy[dim] >= c) ? n.left : n.right; // reverse of above
    if (opposite != INVALID_NODE && limit >= abs((int)xy[dim] - (int)c)) {
      node_distance other_candidate = this->FindNearestRecursive(xy, opposite, level + 1, limit);
      best = SelectNearestNodeDistance(best, other_candidate);
    }
 
    return best;
  }
 
  template <typename Outputter>
  void FindContainedRecursive(CoordT p1[2], CoordT p2[2], size_t node_idx, int level, const Outputter &outputter) const
  {
    /* Dimension index of current level */
    int dim = level % 2;
    /* Node reference */
    const node &n = this->nodes[node_idx];
 
    /* Coordinate of element splitting at this node */
    CoordT ec = TxyFunc()(n.element, dim);
    /* Opposite coordinate of element */
    CoordT oc = TxyFunc()(n.element, 1 - dim);
 
    /* Test if this element is within rectangle */
    if (ec >= p1[dim] && ec < p2[dim] && oc >= p1[1 - dim] && oc < p2[1 - dim]) outputter(n.element);
 
    /* Recurse left if part of rectangle is left of split */
    if (p1[dim] < ec && n.left != INVALID_NODE) this->FindContainedRecursive(p1, p2, n.left, level + 1, outputter);
 
    /* Recurse right if part of rectangle is right of split */
    if (p2[dim] > ec && n.right != INVALID_NODE) this->FindContainedRecursive(p1, p2, n.right, level + 1, outputter);
  }

  size_t CountValue(const T &element, size_t node_idx) const
  {
    if (node_idx == INVALID_NODE) return 0;
    const node &n = this->nodes[node_idx];
    return this->CountValue(element, n.left) + this->CountValue(element, n.right) + ((n.element == element) ? 1 : 0);
  }
 
  void IncrementUnbalanced(size_t amount = 1)
  {
    this->unbalanced += amount;
  }

  bool IsUnbalanced() const
  {
    size_t count = this->Count();
    if (count < MIN_REBALANCE_THRESHOLD) return false;
    return this->unbalanced > count / 4;
  }

  void CheckInvariant(size_t node_idx, int level, CoordT min_x, CoordT max_x, CoordT min_y, CoordT max_y) const
  {
    if (node_idx == INVALID_NODE) return;
 
    const node &n = this->nodes[node_idx];
    CoordT cx = TxyFunc()(n.element, 0);
    CoordT cy = TxyFunc()(n.element, 1);
 
    assert(cx >= min_x);
    assert(cx < max_x);
    assert(cy >= min_y);
    assert(cy < max_y);
 
    if (level % 2 == 0) {
      /* split in dimension 0 = x */
      this->CheckInvariant(n.left,  level + 1, min_x, cx, min_y, max_y);
      this->CheckInvariant(n.right, level + 1, cx, max_x, min_y, max_y);
    } else {
      /* split in dimension 1 = y */
      this->CheckInvariant(n.left,  level + 1, min_x, max_x, min_y, cy);
      this->CheckInvariant(n.right, level + 1, min_x, max_x, cy, max_y);
    }
  }

  void CheckInvariant() const
  {
#ifdef KDTREE_DEBUG
    this->CheckInvariant(this->root, 0, std::numeric_limits<CoordT>::min(), std::numeric_limits<CoordT>::max(), std::numeric_limits<CoordT>::min(), std::numeric_limits<CoordT>::max());
#endif
  }
 
public:
  Kdtree() : root(INVALID_NODE), unbalanced(0) { }

  template <typename It>
  void Build(It begin, It end)
  {
    this->nodes.clear();
    this->free_list.clear();
    this->unbalanced = 0;
    if (begin == end) return;
    this->nodes.reserve(end - begin);
 
    this->root = this->BuildSubtree(begin, end, 0);
    this->CheckInvariant();
  }

  void Clear()
  {
    this->nodes.clear();
    this->free_list.clear();
    this->unbalanced = 0;
    return;
  }

  void Rebuild()
  {
    this->Rebuild(nullptr, nullptr);
  }

  void Insert(const T &element)
  {
    if (this->Count() == 0) {
      this->root = this->AddNode(element);
    } else {
      if (!this->IsUnbalanced() || !this->Rebuild(&element, nullptr)) {
        this->InsertRecursive(element, this->root, 0);
        this->IncrementUnbalanced();
      }
      this->CheckInvariant();
    }
  }

  void Remove(const T &element)
  {
    size_t count = this->Count();
    if (count == 0) return;
    if (!this->IsUnbalanced() || !this->Rebuild(nullptr, &element)) {
      /* If the removed element is the root node, this modifies this->root */
      this->root = this->RemoveRecursive(element, this->root, 0);
      this->IncrementUnbalanced();
    }
    this->CheckInvariant();
  }

  size_t Count() const
  {
    assert(this->free_list.size() <= this->nodes.size());
    return this->nodes.size() - this->free_list.size();
  }

  T FindNearest(CoordT x, CoordT y) const
  {
    assert(this->Count() > 0);
 
    CoordT xy[2] = { x, y };
    return this->FindNearestRecursive(xy, this->root, 0).first;
  }

  template <typename Outputter>
  void FindContained(CoordT x1, CoordT y1, CoordT x2, CoordT y2, const Outputter &outputter) const
  {
    assert(x1 < x2);
    assert(y1 < y2);
 
    if (this->Count() == 0) return;
 
    CoordT p1[2] = { x1, y1 };
    CoordT p2[2] = { x2, y2 };
    this->FindContainedRecursive(p1, p2, this->root, 0, outputter);
  }

  std::vector<T> FindContained(CoordT x1, CoordT y1, CoordT x2, CoordT y2) const
  {
    std::vector<T> result;
    this->FindContained(x1, y1, x2, y2, [&result](T e) {result.push_back(e); });
    return result;
  }
};
 
#endif