DecisionTree Improvements

gtsam/discrete/DecisionTree-inl.h

@@ -635,11 +635,13 @@ namespace gtsam {
       std::function<Y(const X&)> Y_of_X) const {
     using LY = DecisionTree<L, Y>;
 
-    // ugliness below because apparently we can't have templated virtual
-    // functions If leaf, apply unary conversion "op" and create a unique leaf
+    // Ugliness below because apparently we can't have templated virtual
+    // functions.
+    // If leaf, apply unary conversion "op" and create a unique leaf.
     using MXLeaf = typename DecisionTree<M, X>::Leaf;
-    if (auto leaf = boost::dynamic_pointer_cast<const MXLeaf>(f))
+    if (auto leaf = boost::dynamic_pointer_cast<const MXLeaf>(f)) {
       return NodePtr(new Leaf(Y_of_X(leaf->constant())));
+    }
 
     // Check if Choice
     using MXChoice = typename DecisionTree<M, X>::Choice;
@@ -727,6 +729,14 @@ namespace gtsam {
     visit(root_);
   }
 
+  /****************************************************************************/
+  template <typename L, typename Y>
+  size_t DecisionTree<L, Y>::nrLeaves() const {
+    size_t total = 0;
+    visit([&total](const Y& node) { total += 1; });
+    return total;
+  }
+
   /****************************************************************************/
   // fold is just done with a visit
   template <typename L, typename Y>

gtsam/discrete/DecisionTree.h

@@ -262,6 +262,9 @@ namespace gtsam {
     template <typename Func>
     void visitWith(Func f) const;
 
+    /// Return the number of leaves in the tree.
+    size_t nrLeaves() const;
+
     /**
      * @brief Fold a binary function over the tree, returning accumulator.
      *

gtsam/discrete/DecisionTreeFactor.cpp

@@ -156,10 +156,7 @@ namespace gtsam {
   std::vector<std::pair<DiscreteValues, double>> DecisionTreeFactor::enumerate()
       const {
     // Get all possible assignments
-    std::vector<std::pair<Key, size_t>> pairs;
-    for (auto& key : keys()) {
-      pairs.emplace_back(key, cardinalities_.at(key));
-    }
+    std::vector<std::pair<Key, size_t>> pairs = discreteKeys();
     // Reverse to make cartesian product output a more natural ordering.
     std::vector<std::pair<Key, size_t>> rpairs(pairs.rbegin(), pairs.rend());
     const auto assignments = DiscreteValues::CartesianProduct(rpairs);