address review comments

release/4.3a0
Varun Agrawal 2022-03-27 14:35:19 -04:00
parent cd8d5590e6
commit 365473fc5a
1 changed files with 3 additions and 7 deletions

View File

@ -290,10 +290,8 @@ namespace gtsam {
DecisionTreeFactor DecisionTreeFactor::prune(size_t maxNrLeaves) const { DecisionTreeFactor DecisionTreeFactor::prune(size_t maxNrLeaves) const {
const size_t N = maxNrLeaves; const size_t N = maxNrLeaves;
// Let's assume that the structure of the last discrete density will be the // Get the probabilities in the decision tree so we can threshold.
// same as the last continuous
std::vector<double> probabilities; std::vector<double> probabilities;
// number of choices
this->visit([&](const double& prob) { probabilities.emplace_back(prob); }); this->visit([&](const double& prob) { probabilities.emplace_back(prob); });
// The number of probabilities can be lower than max_leaves // The number of probabilities can be lower than max_leaves
@ -318,10 +316,8 @@ namespace gtsam {
}; };
DecisionTree<Key, double> thresholded(*this, thresholdFunc); DecisionTree<Key, double> thresholded(*this, thresholdFunc);
// Create pruned decision tree factor // Create pruned decision tree factor and return.
DecisionTreeFactor prunedDiscreteFactor(this->discreteKeys(), thresholded); return DecisionTreeFactor(this->discreteKeys(), thresholded);
return prunedDiscreteFactor;
} }
/* ************************************************************************ */ /* ************************************************************************ */