KdTree.icpp

Go to the documentation of this file.

 
namespace SourceXtractor {
 
template<typename T, size_t N, size_t S>
class KdTree<T, N, S>::Node {
public:
  virtual std::vector<T> findPointsWithinRadius(Coord coord, double radius) const = 0;
  virtual ~Node() = default;
};
 
template<typename T, size_t N, size_t S>
class KdTree<T, N, S>::Leaf : public KdTree::Node {
public:
  explicit Leaf(const std::vector<T>&& data) : m_data(data) {}
  virtual ~Leaf() = default;
 
  virtual std::vector<T> findPointsWithinRadius(Coord coord, double radius) const {
    std::vector<T> selection;
    for (auto& entry : m_data) {
      double square_dist = 0.0;
      for (size_t i =0; i < N; i++) {
        double delta = Traits::getCoord(entry, i) - coord.coord[i];
        square_dist += delta * delta;
      }
      if (square_dist < radius*radius) {
        selection.push_back(entry);
      }
    }
    return selection;
  }
 
private:
  const std::vector<T> m_data;
};
 
template<typename T, size_t N, size_t S>
class KdTree<T, N, S>::Split : public KdTree::Node  {
public:
  virtual ~Split() = default;
  explicit Split(std::vector<T> data, size_t axis) : m_axis(axis) {
    std::sort(data.begin(), data.end(), [axis](const T& a, const T& b) -> bool {
      return Traits::getCoord(a, axis) < Traits::getCoord(b, axis);
    });
 
    double a = Traits::getCoord(data.at(data.size() / 2 - 1), axis);
    double b = Traits::getCoord(data.at(data.size() / 2), axis);
 
    if (a == b) {
      // avoid a possible rounding issue
      m_split_value = a;
    } else {
      m_split_value = (a + b) / 2.0;
    }
 
    std::vector<T> left(data.begin(), data.begin() + data.size() / 2);
    std::vector<T> right(data.begin() + data.size() / 2, data.end());
 
    if (left.size() > S) {
      m_left_child = std::make_shared<Split>(std::move(left), (axis+1) % N);
    } else {
      m_left_child = std::make_shared<Leaf>(std::move(left));
    }
    if (right.size() > S) {
      m_right_child = std::make_shared<Split>(std::move(right), (axis+1) % N);
    } else {
      m_right_child = std::make_shared<Leaf>(std::move(right));
    }
  }
 
  virtual std::vector<T> findPointsWithinRadius(Coord coord, double radius) const {
    if (coord.coord[m_axis] + radius < m_split_value) {
      return m_left_child->findPointsWithinRadius(coord, radius);
    } else if (coord.coord[m_axis] - radius > m_split_value) {
      return m_right_child->findPointsWithinRadius(coord, radius);
    } else {
      auto left = m_left_child->findPointsWithinRadius(coord, radius);
      auto right = m_right_child->findPointsWithinRadius(coord, radius);
 
      std::vector<T> merge;
      merge.reserve(left.size() + right.size());
      merge.insert(merge.end(), left.begin(), left.end());
      merge.insert(merge.end(), right.begin(), right.end());
 
      return merge;
    }
  }
 
private:
  size_t m_axis;
  double m_split_value;
 
  std::shared_ptr<Node> m_left_child;
  std::shared_ptr<Node> m_right_child;
};
 
template<typename T, size_t N, size_t S>
KdTree<T, N, S>::KdTree(const std::vector<T>& data) {
  if (data.size() > S) {
    m_root = std::make_shared<Split>(data, 0);
  } else {
    std::vector<T> data_copy(data);
    m_root = std::make_shared<Leaf>(std::move(data_copy));
  }
}
 
template<typename T, size_t N, size_t S>
std::vector<T> KdTree<T, N, S>::findPointsWithinRadius(Coord coord, double radius) const {
  return m_root->findPointsWithinRadius(coord, radius);
}
 
}