1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73
| #include <algorithm> #include <cassert> #include <vector>
namespace atcoder {
struct dsu { public: dsu() : _n(0) {} explicit dsu(int n) : _n(n), parent_or_size(n, -1) {} dsu(const dsu&) = default; dsu(dsu&&) = default; dsu& operator=(const dsu&) = default; dsu& operator=(dsu&&) = default;
int merge(int a, int b) { assert(0 <= a && a < _n); assert(0 <= b && b < _n); int x = leader(a), y = leader(b); if (x == y) return x; if (-parent_or_size[x] < -parent_or_size[y]) std::swap(x, y); parent_or_size[x] += parent_or_size[y]; parent_or_size[y] = x; return x; }
bool same(int a, int b) { assert(0 <= a && a < _n); assert(0 <= b && b < _n); return leader(a) == leader(b); }
int leader(int a) { assert(0 <= a && a < _n); int t = a; while (parent_or_size[t] >= 0) t = parent_or_size[t]; while (a != t) { int u = parent_or_size[a]; parent_or_size[a] = t; a = u; } return t; }
int size(int a) { assert(0 <= a && a < _n); return -parent_or_size[leader(a)]; }
std::vector<std::vector<int>> groups() { std::vector<int> leader_buf(_n), group_size(_n); for (int i = 0; i < _n; i++) { leader_buf[i] = leader(i); group_size[leader_buf[i]]++; } std::vector<std::vector<int>> result(_n); for (int i = 0; i < _n; i++) { result[i].reserve(group_size[i]); } for (int i = 0; i < _n; i++) { result[leader_buf[i]].push_back(i); } result.erase(std::remove_if(result.begin(), result.end(), [&](const std::vector<int>& v) { return v.empty(); }), result.end()); return result; }
private: int _n; std::vector<int> parent_or_size; };
}
|