diff --git a/gtsam/hybrid/HybridBayesTree.cpp b/gtsam/hybrid/HybridBayesTree.cpp
index 22777600f..1dc277243 100644
--- a/gtsam/hybrid/HybridBayesTree.cpp
+++ b/gtsam/hybrid/HybridBayesTree.cpp
@@ -47,16 +47,7 @@ DiscreteValues HybridBayesTree::discreteMaxProduct(
     const DiscreteFactorGraph& dfg) const {
   TableFactor product = TableProduct(dfg);
 
-  uint64_t maxIdx = 0;
-  double maxValue = 0.0;
-  Eigen::SparseVector<double> sparseTable = product.sparseTable();
-  for (TableFactor::SparseIt it(sparseTable); it; ++it) {
-    if (it.value() > maxValue) {
-      maxIdx = it.index();
-      maxValue = it.value();
-    }
-  }
-
+  uint64_t maxIdx = TableDistribution(product).argmax();
   DiscreteValues assignment = product.findAssignments(maxIdx);
   return assignment;
 }