dtb.hpp

Go to the documentation of this file.

#ifndef __mlpack_METHODS_EMST_DTB_HPP
#define __mlpack_METHODS_EMST_DTB_HPP

#include "edge_pair.hpp"

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

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

namespace mlpack {
namespace emst  {

class DTBStat
{
 private:
  double maxNeighborDistance;
  int componentMembership;

 public:
  DTBStat();

  template<typename TreeType>
  DTBStat(const TreeType& node);

  double MaxNeighborDistance() const { return maxNeighborDistance; }
  double& MaxNeighborDistance() { return maxNeighborDistance; }

  int ComponentMembership() const { return componentMembership; }
  int& ComponentMembership() { return componentMembership; }

}; // class DTBStat

template<
  typename MetricType = metric::SquaredEuclideanDistance,
  typename TreeType = tree::BinarySpaceTree<bound::HRectBound<2>, DTBStat>
>
class DualTreeBoruvka
{
 private:
  typename TreeType::Mat dataCopy;
  typename TreeType::Mat& data;

  TreeType* tree;
  bool ownTree;

  bool naive;

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

  UnionFind connections;

  std::vector<size_t> oldFromNew;
  arma::Col<size_t> neighborsInComponent;
  arma::Col<size_t> neighborsOutComponent;
  arma::vec neighborsDistances;

  double totalDist;

  MetricType metric;

  // For sorting the edge list after the computation.
  struct SortEdgesHelper
  {
    bool operator()(const EdgePair& pairA, const EdgePair& pairB)
    {
      return (pairA.Distance() < pairB.Distance());
    }
  } SortFun;

 public:
  DualTreeBoruvka(const typename TreeType::Mat& dataset,
                  const bool naive = false,
                  const size_t leafSize = 1,
                  const MetricType metric = MetricType());

  DualTreeBoruvka(TreeType* tree, const typename TreeType::Mat& dataset,
                  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(TreeType* tree);

  void Cleanup();

}; // class DualTreeBoruvka

}; // namespace emst
}; // namespace mlpack

#include "dtb_impl.hpp"

#endif // __mlpack_METHODS_EMST_DTB_HPP