gtsam/gtsam_unstable/base/DSFMap.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 DSFMap.h
 * @date Oct 26, 2013
 * @author Frank Dellaert
 * @brief Allow for arbitrary type in DSF
 */

#pragma once

#include <map>
#include <set>
#include <cstdlib>        // Provides size_t

namespace gtsam {

/**
 * Disjoint set forest using an STL map data structure underneath
 * Uses rank compression and union by rank, iterator version
 * @addtogroup base
 */
template<class KEY>
class DSFMap {

protected:

	/// We store the forest in an STL map, but parents are done with pointers
	struct Entry {
		typename std::map<KEY, Entry>::iterator parent_;
		size_t rank_;
		Entry() {}
	};

    typedef typename std::map<KEY, Entry> Map;
	typedef typename Map::iterator iterator;
	mutable Map entries_;

	/// Given key, find iterator to initial entry
	iterator find__(const KEY& key) const {
		static const Entry empty;
		iterator it = entries_.find(key);
		// if key does not exist, create and return itself
		if (it == entries_.end()) {
			it = entries_.insert(std::make_pair(key, empty)).first;
			it->second.parent_ = it;
			it->second.rank_ = 0;
		}
		return it;
	}

	/// Given iterator to initial entry, find the root Entry
	iterator find_(const iterator& it) const {
		// follow parent pointers until we reach set representative
		iterator& parent = it->second.parent_;
		if (parent != it)
			parent = find_(parent); // not yet, recurse!
		return parent;
	}

	/// Given key, find the root Entry
	inline iterator find_(const KEY& key) const {
		iterator initial = find__(key);
		return find_(initial);
	}

public:

	typedef std::set<KEY> Set;

	/// constructor
	DSFMap() {
	}

	/// Given key, find the representative key for the set in which it lives
	inline KEY find(const KEY& key) const {
		iterator root = find_(key);
		return root->first;
	}

	/// Merge two sets
	void merge(const KEY& x, const KEY& y) {

		// straight from http://en.wikipedia.org/wiki/Disjoint-set_data_structure
		iterator xRoot = find_(x);
		iterator yRoot = find_(y);
		if (xRoot == yRoot)
			return;

		// Merge sets
		if (xRoot->second.rank_ < yRoot->second.rank_)
			xRoot->second.parent_ = yRoot;
		else if (xRoot->second.rank_ > yRoot->second.rank_)
			yRoot->second.parent_ = xRoot;
		else {
			yRoot->second.parent_ = xRoot;
			xRoot->second.rank_ = xRoot->second.rank_ + 1;
		}
	}

	/// return all sets, i.e. a partition of all elements
	std::map<KEY, Set> sets() const {
		std::map<KEY, Set> sets;
		iterator it = entries_.begin();
		for (; it != entries_.end(); it++) {
			iterator root = find_(it);
			sets[root->first].insert(it->first);
		}
		return sets;
	}

};

}