diff options
-rw-r--r-- | source/blender/blenlib/BLI_disjoint_set.hh | 105 | ||||
-rw-r--r-- | source/blender/blenlib/CMakeLists.txt | 1 | ||||
-rw-r--r-- | tests/gtests/blenlib/BLI_disjoint_set_test.cc | 34 | ||||
-rw-r--r-- | tests/gtests/blenlib/CMakeLists.txt | 1 |
4 files changed, 141 insertions, 0 deletions
diff --git a/source/blender/blenlib/BLI_disjoint_set.hh b/source/blender/blenlib/BLI_disjoint_set.hh new file mode 100644 index 00000000000..3b8453669aa --- /dev/null +++ b/source/blender/blenlib/BLI_disjoint_set.hh @@ -0,0 +1,105 @@ +/* + * This program is free software; you can redistribute it and/or + * modify it under the terms of the GNU General Public License + * as published by the Free Software Foundation; either version 2 + * of the License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU General Public License for more details. + * + * You should have received a copy of the GNU General Public License + * along with this program; if not, write to the Free Software Foundation, + * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA. + */ + +#ifndef __BLI_DISJOINT_SET_HH__ +#define __BLI_DISJOINT_SET_HH__ + +/** \file + * \ingroup bli + * + * This implements the disjoint set data structure with path compression and union by rank. + */ + +#include "BLI_array.hh" + +namespace blender { + +class DisjointSet { + private: + Array<uint> parents_; + Array<uint> ranks_; + + public: + /** + * Create a new disjoint set with the given size. Initially, every element is in a separate set. + */ + DisjointSet(uint size) : parents_(size), ranks_(size, 0) + { + for (uint i = 0; i < size; i++) { + parents_[i] = i; + } + } + + /** + * Join the sets containing elements x and y. Nothing happens when they have been in the same set + * before. + */ + void join(uint x, uint y) + { + uint root1 = this->find_root(x); + uint root2 = this->find_root(y); + + /* x and y are in the same set already. */ + if (root1 == root2) { + return; + } + + /* Implement union by rank heuristic. */ + if (ranks_[root1] < ranks_[root2]) { + std::swap(root1, root2); + } + parents_[root2] = root1; + + if (ranks_[root1] == ranks_[root2]) { + ranks_[root1]++; + } + } + + /** + * Return true when x and y are in the same set. + */ + bool in_same_set(uint x, uint y) + { + uint root1 = this->find_root(x); + uint root2 = this->find_root(y); + return root1 == root2; + } + + /** + * Find the element that represents the set containing x currently. + */ + uint find_root(uint x) + { + /* Find root by following parents. */ + uint root = x; + while (parents_[root] != root) { + root = parents_[root]; + } + + /* Compress path. */ + while (parents_[x] != root) { + uint parent = parents_[x]; + parents_[x] = root; + x = parent; + } + + return root; + } +}; + +} // namespace blender + +#endif /* __BLI_DISJOINT_SET_HH__ */ diff --git a/source/blender/blenlib/CMakeLists.txt b/source/blender/blenlib/CMakeLists.txt index b796e893d69..93e5d0a0d79 100644 --- a/source/blender/blenlib/CMakeLists.txt +++ b/source/blender/blenlib/CMakeLists.txt @@ -163,6 +163,7 @@ set(SRC BLI_convexhull_2d.h BLI_delaunay_2d.h BLI_dial_2d.h + BLI_disjoint_set.hh BLI_dlrbTree.h BLI_dot_export.hh BLI_dot_export_attribute_enums.hh diff --git a/tests/gtests/blenlib/BLI_disjoint_set_test.cc b/tests/gtests/blenlib/BLI_disjoint_set_test.cc new file mode 100644 index 00000000000..a46f8612e93 --- /dev/null +++ b/tests/gtests/blenlib/BLI_disjoint_set_test.cc @@ -0,0 +1,34 @@ +#include "BLI_disjoint_set.hh" +#include "BLI_strict_flags.h" + +#include "testing/testing.h" + +namespace blender { + +TEST(disjoint_set, Test) +{ + DisjointSet disjoint_set(6); + EXPECT_FALSE(disjoint_set.in_same_set(1, 2)); + EXPECT_FALSE(disjoint_set.in_same_set(5, 3)); + EXPECT_TRUE(disjoint_set.in_same_set(2, 2)); + EXPECT_EQ(disjoint_set.find_root(3), 3u); + + disjoint_set.join(1, 2); + + EXPECT_TRUE(disjoint_set.in_same_set(1, 2)); + EXPECT_FALSE(disjoint_set.in_same_set(0, 1)); + + disjoint_set.join(3, 4); + + EXPECT_FALSE(disjoint_set.in_same_set(2, 3)); + EXPECT_TRUE(disjoint_set.in_same_set(3, 4)); + + disjoint_set.join(1, 4); + + EXPECT_TRUE(disjoint_set.in_same_set(1, 4)); + EXPECT_TRUE(disjoint_set.in_same_set(1, 3)); + EXPECT_TRUE(disjoint_set.in_same_set(2, 4)); + EXPECT_FALSE(disjoint_set.in_same_set(0, 4)); +} + +} // namespace blender diff --git a/tests/gtests/blenlib/CMakeLists.txt b/tests/gtests/blenlib/CMakeLists.txt index 0831024556b..496fe44234e 100644 --- a/tests/gtests/blenlib/CMakeLists.txt +++ b/tests/gtests/blenlib/CMakeLists.txt @@ -43,6 +43,7 @@ BLENDER_TEST(BLI_array "bf_blenlib") BLENDER_TEST(BLI_array_store "bf_blenlib") BLENDER_TEST(BLI_array_utils "bf_blenlib") BLENDER_TEST(BLI_delaunay_2d "bf_blenlib") +BLENDER_TEST(BLI_disjoint_set "bf_blenlib") BLENDER_TEST(BLI_edgehash "bf_blenlib") BLENDER_TEST(BLI_expr_pylike_eval "bf_blenlib") BLENDER_TEST(BLI_ghash "bf_blenlib") |