gtsam/gtsam_unstable/base/DSF.h

/* ----------------------------------------------------------------------------

 * GTSAM Copyright 2010, Georgia Tech Research Corporation, 
 * Atlanta, Georgia 30332-0415
 * All Rights Reserved
 * Authors: Frank Dellaert, et al. (see THANKS for the full author list)

 * See LICENSE for the license information

 * -------------------------------------------------------------------------- */

/**
 * @file DSF.h
 * @date Mar 26, 2010
 * @author Kai Ni
 * @brief An implementation of Disjoint set forests (see CLR page 446 and up)
 */

#pragma once

#include <iostream>
#include <list>
#include <set>
#include <map>
#include <boost/foreach.hpp>
#include <gtsam_unstable/base/BTree.h>

namespace gtsam {

	/**
	 * Disjoint Set Forest class
	 *
	 * Quoting from CLR: A disjoint-set data structure maintains a collection
	 * S = {S_1,S_2,...} of disjoint dynamic sets. Each set is identified by
	 * a representative, which is some member of the set.
	 *
	 * @addtogroup base
	 */
	template <class KEY>
	class DSF : protected BTree<KEY, KEY> {

	public:
		typedef KEY Label; // label can be different from key, but for now they are same
		typedef DSF<KEY> Self;
		typedef std::set<KEY> Set;
		typedef BTree<KEY, Label> Tree;
		typedef std::pair<KEY, Label> KeyLabel;

		// constructor
		DSF() : Tree() { }

		// constructor
		DSF(const Tree& tree) : Tree(tree) {}

		// constructor with a list of unconnected keys
		DSF(const std::list<KEY>& keys) : Tree() { BOOST_FOREACH(const KEY& key, keys) *this = this->add(key, key); }

		// constructor with a set of unconnected keys
		DSF(const std::set<KEY>& keys) : Tree() { BOOST_FOREACH(const KEY& key, keys) *this = this->add(key, key); }

		// create a new singleton, does nothing if already exists
		Self makeSet(const KEY& key) const { if (this->mem(key)) return *this; else return this->add(key, key); }

		// find the label of the set in which {key} lives
		Label findSet(const KEY& key) const {
			KEY parent = this->find(key);
			return parent == key ? key : findSet(parent); }

		// return a new DSF where x and y are in the same set. Kai: the caml implementation is not const, and I followed
		Self makeUnion(const KEY& key1, const KEY& key2) { return this->add(findSet_(key2), findSet_(key1));	}

		// the in-place version of makeUnion
		void makeUnionInPlace(const KEY& key1, const KEY& key2) { *this = this->add(findSet_(key2), findSet_(key1)); }

		// create a new singleton with two connected keys
		Self makePair(const KEY& key1, const KEY& key2) const { return makeSet(key1).makeSet(key2).makeUnion(key1, key2); }

		// create a new singleton with a list of fully connected keys
		Self makeList(const std::list<KEY>& keys) const {
			Self t = *this;
			BOOST_FOREACH(const KEY& key, keys)
				t = t.makePair(key, keys.front());
			return t;
		}

		// return a dsf in which all find_set operations will be O(1) due to path compression.
		DSF flatten() const {
			DSF t = *this;
			BOOST_FOREACH(const KeyLabel& pair, (Tree)t)
				t.findSet_(pair.first);
			return t;
		}

		// maps f over all keys, must be invertible
		DSF map(boost::function<KEY(const KEY&)> func) const {
			DSF t;
			BOOST_FOREACH(const KeyLabel& pair, (Tree)*this)
				t = t.add(func(pair.first), func(pair.second));
			return t;
		}

		// return the number of sets
		size_t numSets() const {
			size_t num = 0;
			BOOST_FOREACH(const KeyLabel& pair, (Tree)*this)
				if (pair.first == pair.second) num++;
			return num;
		}

		// return the numer of keys
		size_t size() const { return Tree::size(); }

		// return all sets, i.e. a partition of all elements
		std::map<Label, Set> sets() const {
			std::map<Label, Set> sets;
			BOOST_FOREACH(const KeyLabel& pair, (Tree)*this)
				sets[findSet(pair.second)].insert(pair.first);
			return sets;
		}

		// return a partition of the given elements {keys}
		std::map<Label, Set> partition(const std::list<KEY>& keys) const {
			std::map<Label, Set> partitions;
			BOOST_FOREACH(const KEY& key, keys)
				partitions[findSet(key)].insert(key);
			return partitions;
		}

		// get the nodes in the tree with the given label
		Set set(const Label& label) const {
			Set set;
			BOOST_FOREACH(const KeyLabel& pair, (Tree)*this) {
				if (pair.second == label || findSet(pair.second) == label)
					set.insert(pair.first);
			}
			return set;
		}

		/** equality */
		bool operator==(const Self& t) const { return (Tree)*this == (Tree)t;	}

		/** inequality */
		bool operator!=(const Self& t) const { return (Tree)*this != (Tree)t;	}

		// print the object
		void print(const std::string& name = "DSF") const {
			std::cout << name << std::endl;
			BOOST_FOREACH(const KeyLabel& pair, (Tree)*this)
				std::cout << (std::string)pair.first << " " << (std::string)pair.second << std::endl;
		}

	protected:

		/**
		 * same as findSet except with path compression: After we have traversed the path to
		 * the root, each parent pointer is made to directly point to it
		 */
		KEY findSet_(const KEY& key) {
			KEY parent = this->find(key);
			if (parent == key)
				return parent;
			else {
				KEY label = findSet_(parent);
				*this = this->add(key, label);
				return label;
			}
		}

	};

	// shortcuts
	typedef DSF<int> DSFInt;

} // namespace gtsam