dtb.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_EMST_DTB_HPP
#define MLPACK_METHODS_EMST_DTB_HPP

#include "dtb_stat.hpp"
#include "edge_pair.hpp"

#include <mlpack/prereqs.hpp>
#include <mlpack/core/metrics/lmetric.hpp>

#include <mlpack/core/tree/binary_space_tree.hpp>

namespace mlpack {
namespace emst  {

template<
    typename MetricType = metric::EuclideanDistance,
    typename MatType = arma::mat,
    template<typename TreeMetricType,
             typename TreeStatType,
             typename TreeMatType> class TreeType = tree::KDTree
>
class DualTreeBoruvka
{
 public:
  typedef TreeType<MetricType, DTBStat, MatType> Tree;

 private:
  std::vector<size_t> oldFromNew;
  Tree* tree;
  const MatType& data;
  bool ownTree;

  bool naive;

  std::vector<EdgePair> edges; // We must use vector with non-numerical types.

  UnionFind connections;

  arma::Col<size_t> neighborsInComponent;
  arma::Col<size_t> neighborsOutComponent;
  arma::vec neighborsDistances;

  double totalDist;

  MetricType metric;

  struct SortEdgesHelper
  {
    bool operator()(const EdgePair& pairA, const EdgePair& pairB)
    {
      return (pairA.Distance() < pairB.Distance());
    }
  } SortFun;

 public:
  DualTreeBoruvka(const MatType& dataset,
                  const bool naive = false,
                  const MetricType metric = MetricType());

  DualTreeBoruvka(Tree* tree,
                  const MetricType metric = MetricType());

  ~DualTreeBoruvka();

  void ComputeMST(arma::mat& results);

 private:
  void AddEdge(const size_t e1, const size_t e2, const double distance);

  void AddAllEdges();

  void EmitResults(arma::mat& results);

  void CleanupHelper(Tree* tree);

  void Cleanup();

}; // class DualTreeBoruvka

} // namespace emst
} // namespace mlpack

#include "dtb_impl.hpp"

#endif // MLPACK_METHODS_EMST_DTB_HPP