diff --git a/gtsam/discrete/DecisionTreeFactor.cpp b/gtsam/discrete/DecisionTreeFactor.cpp
index 19deccd78..0396dc952 100644
--- a/gtsam/discrete/DecisionTreeFactor.cpp
+++ b/gtsam/discrete/DecisionTreeFactor.cpp
@@ -247,7 +247,7 @@ namespace gtsam {
   /* ************************************************************************ */
   std::vector<double> DecisionTreeFactor::probabilities() const {
     // Set of all keys
-    std::set<Key> allKeys(keys().begin(), keys().end());
+    KeySet allKeys(keys().begin(), keys().end());
 
     std::vector<double> probs;
 
@@ -260,7 +260,7 @@ namespace gtsam {
      */
     auto op = [&](const Assignment<Key>& a, double p) {
       // Get all the keys in the current assignment
-      std::set<Key> assignment_keys;
+      KeySet assignment_keys;
       for (auto&& [k, _] : a) {
         assignment_keys.insert(k);
       }
@@ -453,18 +453,18 @@ namespace gtsam {
   /* ************************************************************************ */
   double DecisionTreeFactor::computeThreshold(const size_t N) const {
     // Set of all keys
-    std::set<Key> allKeys = this->labels();
+    KeySet allKeys = this->labels();
     MinHeap min_heap;
 
     auto op = [&](const Assignment<Key>& a, double p) {
       // Get all the keys in the current assignment
-      std::set<Key> assignment_keys;
+      KeySet assignment_keys;
       for (auto&& [k, _] : a) {
         assignment_keys.insert(k);
       }
 
       // Find the keys missing in the assignment
-      std::vector<Key> diff;
+      KeyVector diff;
       std::set_difference(allKeys.begin(), allKeys.end(),
                           assignment_keys.begin(), assignment_keys.end(),
                           std::back_inserter(diff));