Created
December 22, 2020 16:09
-
-
Save thorstel/aa9e03516bd57674f806b77933badbd3 to your computer and use it in GitHub Desktop.
Disjoint Sets / Union-Find
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #ifndef DISJOINT_SET_HH_ | |
| #define DISJOINT_SET_HH_ | |
| #include <map> | |
| template<typename T> | |
| class disjoint_set | |
| { | |
| struct elem_t | |
| { | |
| T id; | |
| int rank {0}; | |
| size_t size {0}; | |
| }; | |
| std::map<T, elem_t> forest; | |
| size_t num_sets {0}; | |
| public: | |
| /// Finds the set which contains elem (or creates a new one). | |
| T find(const T& elem) { return find_rec(elem).id; } | |
| /// Checks if elem is contained in any set. | |
| bool contains(const T& elem) const { return forest.find(elem) != std::end(forest); } | |
| /// Returns the size of the set which contains elem (or creates a new one with size 1). | |
| size_t size_of(const T& elem) { return find_rec(elem).size; } | |
| /// Returns the number of sets. | |
| size_t size() const { return num_sets; } | |
| /// Joins the specified sets. | |
| void make_union(const T& elem1, const T& elem2) | |
| { | |
| auto& set1 = find_rec(elem1); | |
| auto& set2 = find_rec(elem2); | |
| if (set1.id == set2.id) { | |
| return; | |
| } | |
| if (set1.rank < set2.rank) { | |
| set1.id = set2.id; | |
| set2.size += set1.size; | |
| } else { | |
| set2.id = set1.id; | |
| set1.size += set2.size; | |
| if (set1.rank == set2.rank) { ++set1.rank; } | |
| } | |
| --num_sets; | |
| } | |
| private: | |
| /// Recursive implementation of union-find. | |
| elem_t& find_rec(const T& elem) | |
| { | |
| auto& s = forest[elem]; | |
| if (s.size == 0) { | |
| ++num_sets; | |
| s.id = elem; | |
| s.size = 1; | |
| return s; | |
| } else if (s.id == elem) { | |
| return s; | |
| } else { | |
| auto& parent = find_rec(s.id); | |
| s.id = parent.id; | |
| return parent; | |
| } | |
| } | |
| }; | |
| #endif /* DISJOINT_SET_HH_ */ |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment