update restrict to remove restricted keys
parent
b7d0931fd9
commit
1b82faeb00
|
@ -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);
|
||||
}
|
||||
|
||||
/* ************************************************************************ */
|
||||
|
|
Loading…
Reference in New Issue