Theory

• Algorithms, 4th edition by Robert Sedgewick, Kevin Wayne
  • 1.5   Case Study: Union-Find
  • Union Find on Cubits.ai

Implementation in C++

#include <bits/stdc++.h>

using namespace std;

class DSU {
    private:
        vector<int> parent;
        vector<int> size;

        int root(int n) {
            while (parent[n] != n) {
                // Path compression
                parent[n] = parent[parent[n]];
                n = parent[n];
            }
            return n;
        }

    public:
        DSU(int n) {
            parent.resize(n);
            std::iota(begin(parent), end(parent), 0);
            size.resize(n);
            std::fill(begin(size), end(size), 1);
        }

        void join(int p, int q) {
            int x = root(p);
            int y = root(q);
            if (x == y) {
                return;
            }

            if (size[x] < size[y]) {
                size[y] += size[x];
                parent[x] = y;
            } else {
                size[x] += size[y];
                parent[y] = x;
            }
        }

        bool connected(int p, int q) {
            return root(p) == root(q);
        }
};

Practice

• https://atcoder.jp/contests/practice2/tasks/practice2_a