add nrAssignments method for DecisionTree

release/4.3a0
Varun Agrawal 2023-03-26 16:02:36 -04:00
parent 005c7d4e2d
commit e114e9f6d2
2 changed files with 21 additions and 1 deletions

View File

@ -649,8 +649,9 @@ namespace gtsam {
throw std::invalid_argument("DecisionTree::create invalid argument");
}
auto choice = std::make_shared<Choice>(begin->first, endY - beginY);
for (ValueIt y = beginY; y != endY; y++)
for (ValueIt y = beginY; y != endY; y++) {
choice->push_back(NodePtr(new Leaf(*y)));
}
return Choice::Unique(choice);
}
@ -827,6 +828,16 @@ namespace gtsam {
return total;
}
/****************************************************************************/
template <typename L, typename Y>
size_t DecisionTree<L, Y>::nrAssignments() const {
size_t n = 0;
this->visitLeaf([&n](const DecisionTree<L, Y>::Leaf& leaf) {
n += leaf.nrAssignments();
});
return n;
}
/****************************************************************************/
// fold is just done with a visit
template <typename L, typename Y>

View File

@ -299,6 +299,15 @@ namespace gtsam {
/// Return the number of leaves in the tree.
size_t nrLeaves() const;
/**
* @brief Return the number of total leaf assignments.
* This includes counts removed from implicit pruning hence,
* it will always be >= nrLeaves().
*
* @return size_t
*/
size_t nrAssignments() const;
/**
* @brief Fold a binary function over the tree, returning accumulator.
*