update restrict to remove restricted keys

release/4.3a0
Varun Agrawal 2025-04-11 15:04:26 -04:00
parent b7d0931fd9
commit 1b82faeb00
1 changed files with 10 additions and 2 deletions

View File

@ -548,8 +548,16 @@ namespace gtsam {
DiscreteFactor::shared_ptr DecisionTreeFactor::restrict(
const DiscreteValues& assignment) const {
ADT restricted_tree = ADT::restrict(assignment);
return std::make_shared<DecisionTreeFactor>(this->discreteKeys(),
restricted_tree);
DiscreteKeys restricted_keys = this->discreteKeys();
for (auto&& kv : assignment) {
Key key = kv.first;
// Remove the key from the keys list
restricted_keys.erase(
std::remove_if(restricted_keys.begin(), restricted_keys.end(),
[key](const DiscreteKey& k) { return k.first == key; }),
restricted_keys.end());
}
return std::make_shared<DecisionTreeFactor>(restricted_keys, restricted_tree);
}
/* ************************************************************************ */