algo/ds/dsu.h

• View this file on GitHub
• Last update: 2026-09-28 20:05:44-04:00

Depends on

• algo/common.h

Verified with

• verify/ds/dsu.test.cpp

Code

#pragma once
#include "algo/common.h"

namespace algo::ds {

// Disjoint set union.
//   dsu d(n);
//   d.unite(a, b);    // false if already joined
//   d.is_same(a, b);
//   d.size(a);        // size of a's set
template <bool union_by_size = true, bool path_compression = true>
struct dsu {
    dsu(int n) : e(std::vector<int>(n, -1)) {
    }
    int get(int x) {
        if (e[x] < 0) return x;
        if (path_compression) return e[x] = get(e[x]);
        return get(e[x]);
    }
    bool is_same(int a, int b) {
        return get(a) == get(b);
    }
    int size(int x) {
        return -e[get(x)];
    }
    bool unite(int x, int y) {
        x = get(x), y = get(y);
        if (x == y) return false;
        if (union_by_size && e[x] > e[y]) std::swap(x, y);
        e[x] += e[y];
        e[y] = x;
        return true;
    }
    friend std::ostream &operator<<(std::ostream &os, dsu s) {
        os << "[";
        bool first = true;
        for (int i = 0; i < (int)s.e.size(); i++) {
            if (s.get(i) == i) {
                if (!first) os << ", ";
                first = false;
                os << "[" << i;
                for (int j = 0; j < (int)s.e.size(); j++) {
                    if (j != i && s.get(j) == i) {
                        os << ", " << j;
                    }
                }
                os << "]";
            }
        }
        return os << "]";
    }

private:
    // Root: -(set size). Otherwise: parent.
    std::vector<int> e;
};

} // namespace algo::ds

#line 2 "algo/common.h"
#ifndef PREPROCESS
#include <bits/stdc++.h>
#include <cassert>
#endif

// Declared here so `using namespace algo;` works with no other includes.
namespace algo {}
#line 3 "algo/ds/dsu.h"

namespace algo::ds {

// Disjoint set union.
//   dsu d(n);
//   d.unite(a, b);    // false if already joined
//   d.is_same(a, b);
//   d.size(a);        // size of a's set
template <bool union_by_size = true, bool path_compression = true>
struct dsu {
    dsu(int n) : e(std::vector<int>(n, -1)) {
    }
    int get(int x) {
        if (e[x] < 0) return x;
        if (path_compression) return e[x] = get(e[x]);
        return get(e[x]);
    }
    bool is_same(int a, int b) {
        return get(a) == get(b);
    }
    int size(int x) {
        return -e[get(x)];
    }
    bool unite(int x, int y) {
        x = get(x), y = get(y);
        if (x == y) return false;
        if (union_by_size && e[x] > e[y]) std::swap(x, y);
        e[x] += e[y];
        e[y] = x;
        return true;
    }
    friend std::ostream &operator<<(std::ostream &os, dsu s) {
        os << "[";
        bool first = true;
        for (int i = 0; i < (int)s.e.size(); i++) {
            if (s.get(i) == i) {
                if (!first) os << ", ";
                first = false;
                os << "[" << i;
                for (int j = 0; j < (int)s.e.size(); j++) {
                    if (j != i && s.get(j) == i) {
                        os << ", " << j;
                    }
                }
                os << "]";
            }
        }
        return os << "]";
    }

private:
    // Root: -(set size). Otherwise: parent.
    std::vector<int> e;
};

} // namespace algo::ds

Back to top page