85 lines
1.6 KiB
C++
85 lines
1.6 KiB
C++
#pragma once
|
|
|
|
#include <memory>
|
|
|
|
template <class uintXX_t = uint32_t,
|
|
class allocator = std::allocator<uintXX_t>>
|
|
class UnionFind
|
|
{
|
|
public:
|
|
UnionFind(uintXX_t n) : N(n), n(n)
|
|
{
|
|
count = alloc.allocate(n);
|
|
parent = alloc.allocate(n);
|
|
|
|
for(auto i = 0; i < n; ++i)
|
|
count[i] = 1, parent[i] = i;
|
|
}
|
|
|
|
~UnionFind()
|
|
{
|
|
alloc.deallocate(count, N);
|
|
alloc.deallocate(parent, N);
|
|
}
|
|
|
|
// this is O(lg* n)
|
|
void connect(uintXX_t p, uintXX_t q)
|
|
{
|
|
auto rp = root(p);
|
|
auto rq = root(q);
|
|
|
|
// if roots are equal, we don't have to do anything
|
|
if(rp == rq)
|
|
return;
|
|
|
|
// merge the smaller subtree to the root of the larger subtree
|
|
if(count[rp] < count[rq])
|
|
parent[rp] = rq, count[rp] += count[rp];
|
|
else
|
|
parent[rq] = rp, count[rp] += count[rq];
|
|
|
|
// update the number of groups
|
|
n--;
|
|
}
|
|
|
|
// O(lg* n)
|
|
bool find(uintXX_t p, uintXX_t q)
|
|
{
|
|
return root(p) == root(q);
|
|
}
|
|
|
|
// O(lg* n)
|
|
uintXX_t root(uintXX_t p)
|
|
{
|
|
auto r = p;
|
|
auto newp = p;
|
|
|
|
// find the node connected to itself, that's the root
|
|
while(parent[r] != r)
|
|
r = parent[r];
|
|
|
|
// do some path compression to enable faster searches
|
|
while(p != r)
|
|
newp = parent[p], parent[p] = r, p = newp;
|
|
|
|
return r;
|
|
}
|
|
|
|
uintXX_t size() const
|
|
{
|
|
return n;
|
|
}
|
|
|
|
private:
|
|
allocator alloc;
|
|
|
|
const uintXX_t N;
|
|
uintXX_t n;
|
|
|
|
// array of subtree counts
|
|
uintXX_t* count;
|
|
|
|
// array of tree indices
|
|
uintXX_t* parent;
|
|
};
|