diff --git a/src/mlpack/core/tree/binary_space_tree_impl.hpp b/src/mlpack/core/tree/binary_space_tree_impl.hpp index ba98b29a9b..aa0c8b6479 100644 --- a/src/mlpack/core/tree/binary_space_tree_impl.hpp +++ b/src/mlpack/core/tree/binary_space_tree_impl.hpp @@ -99,6 +99,12 @@ BinarySpaceTree::BinarySpaceTree( { // Perform the actual splitting. SplitNode(data); + + // Create the statistic depending on if we are a leaf or not. + if (IsLeaf()) + stat = StatisticType(data, begin, count); + else + stat = StatisticType(data, begin, count, left->Stat(), right->Stat()); } template @@ -122,6 +128,12 @@ BinarySpaceTree::BinarySpaceTree( // Perform the actual splitting. SplitNode(data, oldFromNew); + + // Create the statistic depending on if we are a leaf or not. + if (IsLeaf()) + stat = StatisticType(data, begin, count); + else + stat = StatisticType(data, begin, count, left->Stat(), right->Stat()); } template @@ -151,6 +163,12 @@ BinarySpaceTree::BinarySpaceTree( newFromOld.resize(data.n_cols); for (size_t i = 0; i < data.n_cols; i++) newFromOld[oldFromNew[i]] = i; + + // Create the statistic depending on if we are a leaf or not. + if (IsLeaf()) + stat = StatisticType(data, begin, count); + else + stat = StatisticType(data, begin, count, left->Stat(), right->Stat()); } template diff --git a/src/mlpack/core/tree/statistic.hpp b/src/mlpack/core/tree/statistic.hpp index 35fc0e5207..1235f24168 100644 --- a/src/mlpack/core/tree/statistic.hpp +++ b/src/mlpack/core/tree/statistic.hpp @@ -15,8 +15,6 @@ namespace tree { /** * Empty statistic if you are not interested in storing statistics in your * tree. Use this as a template for your own. - * - * @experimental */ class EmptyStatistic { @@ -25,17 +23,35 @@ class EmptyStatistic ~EmptyStatistic() {} /** - * Initializes by taking statistics on raw data. + * This constructor is called when a leaf is created. + * + * @param dataset Matrix that the tree is being built on. + * @param begin Starting index corresponding to this leaf. + * @param count Number of points held in this leaf. */ - void Init(const arma::mat& dataset, size_t start, size_t count) { } + template + EmptyStatistic(const MatType& dataset, + const size_t begin, + const size_t count) + { } /** - * Initializes by combining statistics of two partitions. - * + * This constructor is called when a non-leaf node is created. * This lets you build fast bottom-up statistics when building trees. + * + * @param dataset Matrix that the tree is being built on. + * @param begin Starting index corresponding to this leaf. + * @param count Number of points held in this leaf. + * @param leftStat EmptyStatistic object of the left child node. + * @param rightStat EmptyStatistic object of the right child node. */ - void Init(const arma::mat& dataset, size_t start, size_t count, - const EmptyStatistic& left_stat, const EmptyStatistic& right_stat) { } + template + EmptyStatistic(const MatType& dataset, + const size_t start, + const size_t count, + const EmptyStatistic& leftStat, + const EmptyStatistic& rightStat) + { } }; }; // namespace tree