From 0142cc3d3f427652d3fbbc07d31fd333dd3ae52c Mon Sep 17 00:00:00 2001
From: Kai Ni <nikai97@gmail.com>
Date: Fri, 9 Apr 2010 23:36:13 +0000
Subject: [PATCH] add nrKeys and set

---
 cpp/DSF.h             |  8 ++++++++
 cpp/FactorGraph-inl.h |  3 +--
 cpp/FactorGraph.h     |  3 +++
 cpp/testDSF.cpp       | 14 ++++++++++++++
 4 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/cpp/DSF.h b/cpp/DSF.h
index c66ba509b..82b9bb4d0 100644
--- a/cpp/DSF.h
+++ b/cpp/DSF.h
@@ -105,6 +105,14 @@ namespace gtsam {
 			return partitions;
 		}
 
+		// get the nodes in the given tree
+		Set set(const Label& label) {
+			Set set;
+			BOOST_FOREACH(const KeyLabel& pair, (Tree)*this)
+				if (pair.second==label) set.insert(pair.first);
+			return set;
+		}
+
 		/** equality */
 		bool operator==(const Self& t) const { return (Tree)*this == (Tree)t;	}
 
diff --git a/cpp/FactorGraph-inl.h b/cpp/FactorGraph-inl.h
index 3633f5862..04e7d630a 100644
--- a/cpp/FactorGraph-inl.h
+++ b/cpp/FactorGraph-inl.h
@@ -407,8 +407,7 @@ std::pair<FactorGraph<Factor>, FactorGraph<Factor> > FactorGraph<Factor>::splitM
 	DSF<Symbol> dsf(keys());
 
 	//	while G is nonempty and T is not yet spanning
-	size_t m = nrFactors();
-	for (size_t i=0;i<m;i++) {
+	for (size_t i=0;i<size();i++) {
 		const sharedFactor& f = factors_[i];
 
 		// retrieve the labels of all the keys
diff --git a/cpp/FactorGraph.h b/cpp/FactorGraph.h
index 00d9f9c0e..fb0193c75 100644
--- a/cpp/FactorGraph.h
+++ b/cpp/FactorGraph.h
@@ -92,6 +92,9 @@ namespace gtsam {
 		/** return keys in some random order */
 		Ordering keys() const;
 
+		/** return the number of the keys */
+		inline size_t nrKeys() const {return indices_.size(); };
+
 		/** Check whether a factor with this variable exists */
 		bool involves(const Symbol& key) const {
 			return !(indices_.find(key)==indices_.end());
diff --git a/cpp/testDSF.cpp b/cpp/testDSF.cpp
index 10965f32d..7f4f2ce1f 100644
--- a/cpp/testDSF.cpp
+++ b/cpp/testDSF.cpp
@@ -181,6 +181,20 @@ TEST(DSF, partition3) {
 	CHECK(expected == partitions[dsf.findSet(5)]);
 }
 
+/* ************************************************************************* */
+TEST(DSF, set) {
+	DSFInt dsf;
+	dsf = dsf.makeSet(5);
+	dsf = dsf.makeSet(6);
+	dsf = dsf.makeSet(7);
+	dsf = dsf.makeUnion(5,6);
+	set<int> set = dsf.set(5);
+	LONGS_EQUAL(2, set.size());
+
+	std::set<int> expected; expected += 5, 6;
+	CHECK(expected == set);
+}
+
 /* ************************************************************************* */
 int func(const int& a) { return a + 10; }
 TEST(DSF, map) {