From 1b82faeb0009e65b19f61e48384a2ac22421b6ed Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Fri, 11 Apr 2025 15:04:26 -0400
Subject: [PATCH] update restrict to remove restricted keys

---
 gtsam/discrete/DecisionTreeFactor.cpp | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/gtsam/discrete/DecisionTreeFactor.cpp b/gtsam/discrete/DecisionTreeFactor.cpp
index ebc02c1b5..1ba52642c 100644
--- a/gtsam/discrete/DecisionTreeFactor.cpp
+++ b/gtsam/discrete/DecisionTreeFactor.cpp
@@ -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);
 }
 
   /* ************************************************************************ */