diff --git a/libs/util/include/psemek/util/dsu.hpp b/libs/util/include/psemek/util/dsu.hpp new file mode 100644 index 00000000..63765875 --- /dev/null +++ b/libs/util/include/psemek/util/dsu.hpp @@ -0,0 +1,54 @@ +#pragma once + +#include +#include + +namespace psemek::util +{ + + template + struct dsu + { + dsu() = default; + dsu(Index count); + + void combine(Index x, Index y); + Index find(Index x); + + private: + std::vector parent_; + std::vector size_; + }; + + template + dsu::dsu(Index count) + : parent_(count) + , size_(count, 1) + { + for (Index i = 0; i < count; ++i) + parent_[i] = i; + } + + template + void dsu::combine(Index x, Index y) + { + x = find(x); + y = find(y); + if (x != y) + { + if (size_[x] < size_[y]) + std::swap(x, y); + parent_[y] = x; + size_[x] += size_[y]; + } + } + + template + Index dsu::find(Index x) + { + if (parent_[x] == x) + return x; + return parent_[x] = find(parent_[x]); + } + +}