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