diff --git a/gtsam/discrete/DiscreteFactorGraph.cpp b/gtsam/discrete/DiscreteFactorGraph.cpp
index c1248c60b..cb7c16595 100644
--- a/gtsam/discrete/DiscreteFactorGraph.cpp
+++ b/gtsam/discrete/DiscreteFactorGraph.cpp
@@ -43,11 +43,25 @@ namespace gtsam {
   /* ************************************************************************* */
   KeySet DiscreteFactorGraph::keys() const {
     KeySet keys;
-    for(const sharedFactor& factor: *this)
-    if (factor) keys.insert(factor->begin(), factor->end());
+    for (const sharedFactor& factor : *this) {
+      if (factor) keys.insert(factor->begin(), factor->end());
+    }
     return keys;
   }
 
+  /* ************************************************************************* */
+  DiscreteKeys DiscreteFactorGraph::discreteKeys() const {
+    DiscreteKeys result;
+    for (auto&& factor : *this) {
+      if (auto p = boost::dynamic_pointer_cast<DecisionTreeFactor>(factor)) {
+        DiscreteKeys factor_keys = p->discreteKeys();
+        result.insert(result.end(), factor_keys.begin(), factor_keys.end());
+      }
+    }
+
+    return result;
+  }
+
   /* ************************************************************************* */
   DecisionTreeFactor DiscreteFactorGraph::product() const {
     DecisionTreeFactor result;
diff --git a/gtsam/discrete/DiscreteFactorGraph.h b/gtsam/discrete/DiscreteFactorGraph.h
index 1da840eb8..baa0213d5 100644
--- a/gtsam/discrete/DiscreteFactorGraph.h
+++ b/gtsam/discrete/DiscreteFactorGraph.h
@@ -114,6 +114,9 @@ class GTSAM_EXPORT DiscreteFactorGraph
   /** Return the set of variables involved in the factors (set union) */
   KeySet keys() const;
 
+  /// Return the DiscreteKeys in this factor graph.
+  DiscreteKeys discreteKeys() const;
+
   /** return product of all factors as a single factor */
   DecisionTreeFactor product() const;