From 69b5e7d5275f28392dff6f844514b6ee92454bc8 Mon Sep 17 00:00:00 2001
From: Varun Agrawal <varagrawal@gmail.com>
Date: Sat, 4 Jan 2025 17:10:43 -0500
Subject: [PATCH] return DiscreteValues directly

---
 gtsam/discrete/TableDistribution.cpp | 4 ++--
 gtsam/discrete/TableDistribution.h   | 2 +-
 gtsam/hybrid/HybridBayesTree.cpp     | 3 +--
 3 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/gtsam/discrete/TableDistribution.cpp b/gtsam/discrete/TableDistribution.cpp
index 241320649..aa639c126 100644
--- a/gtsam/discrete/TableDistribution.cpp
+++ b/gtsam/discrete/TableDistribution.cpp
@@ -119,7 +119,7 @@ DiscreteFactor::shared_ptr TableDistribution::max(const Ordering& keys) const {
 }
 
 /* ************************************************************************ */
-uint64_t TableDistribution::argmax() const {
+DiscreteValues TableDistribution::argmax() const {
   uint64_t maxIdx = 0;
   double maxValue = 0.0;
 
@@ -132,7 +132,7 @@ uint64_t TableDistribution::argmax() const {
     }
   }
 
-  return maxIdx;
+  return table_.findAssignments(maxIdx);
 }
 
 /* ****************************************************************************/
diff --git a/gtsam/discrete/TableDistribution.h b/gtsam/discrete/TableDistribution.h
index 662602c77..65e895a85 100644
--- a/gtsam/discrete/TableDistribution.h
+++ b/gtsam/discrete/TableDistribution.h
@@ -145,7 +145,7 @@ class GTSAM_EXPORT TableDistribution : public DiscreteConditional {
    *
    * @return maximizing assignment for the variables.
    */
-  uint64_t argmax() const;
+  DiscreteValues argmax() const;
 
   /// @}
   /// @name Advanced Interface
diff --git a/gtsam/hybrid/HybridBayesTree.cpp b/gtsam/hybrid/HybridBayesTree.cpp
index 1dc277243..0df46f262 100644
--- a/gtsam/hybrid/HybridBayesTree.cpp
+++ b/gtsam/hybrid/HybridBayesTree.cpp
@@ -47,8 +47,7 @@ DiscreteValues HybridBayesTree::discreteMaxProduct(
     const DiscreteFactorGraph& dfg) const {
   TableFactor product = TableProduct(dfg);
 
-  uint64_t maxIdx = TableDistribution(product).argmax();
-  DiscreteValues assignment = product.findAssignments(maxIdx);
+  DiscreteValues assignment = TableDistribution(product).argmax();
   return assignment;
 }