Пособие/«Система непересекающихся множеств»

#include <iostream>
#include <vector>
#include <string>

using namespace std;

vector<int> p, sz; // создаём вектора p(массив ссылок), sz(массив размеров множеств)

int get_p(int v) { // get_p(v) возвращает номер множества, которому принадлежит v-ый элемент
    if (p[v] == v) { // Если элемент ссылается на самого себя, то это корень, следовательно это искомый ответ
        return v; // возвращаем искомый ответ
    }
    // одновременно находим корень множества (оно в виде дерева), к которому принадлежит предок v и напрямую подвешиваем к этому корню v
    return p[v] = get_p(p[v]);
}

void unite(int a, int b) { // функция, объединяющая 2 множества
    a = get_p(a); // находим корень множества, к которому принадлежит a
    b = get_p(b); // находим корень множества, к которому принадлежит b
    // Если размер множества, в котором корень a, больше размера множества, в котором корень b, то
    // меняем значения в a и b
    if (sz[a] > sz[b]) {
        swap(a, b);
    }
    sz[b] += sz[a]; // увеличиваем размер множества с корнем b на размер множества с корнем a
    p[a] = b; // подвешиваем a к b
}

void solve() {
    string query; // query - это очередной запрос
    while (cin >> query) { // Если какой-то запрос считывается, то продолжаем работу программы
        if (query == "RESET") { // если запрос - это "RESET", то сбрасываем все массивы
            int n;
            cin >> n; // считываем кол-во элементов
            p.assign(n, 0); // меняем размер массива p; неважно, чем мы его заполним, так как эти значения мы всё равно присвоим циклом
            sz.assign(n, 1); // меняем размер массива sz, заполняем его 1-ами, так как изначально i-ое множество состоит из одного элемента (i-го элемента)
            for (int i = 0; i < n; ++i) p[i] = i; // так как изначально i-ое множество состоит только из i-го элемента, то i-ый элемент является корнем, а значит ссылается на самого себя
            cout << "RESET DONE\n"; // Вывод этой фразы просто прописан в условии задачи
        }
        if (query == "JOIN") { // если запрос - это "JOIN", то нужно объединить множество, в котором содержится j-ый элемент и множество, в котором содержится k-ый элемент
            int j, k;
            cin >> j >> k; // считываем j и k
            if (get_p(j) == get_p(k)) { // Если номера множеств, к которым принадлежат j и k совпадают, то они уже лежат в одном множестве
                cout << "ALREADY " << j << ' ' << k << '\n'; // Вывод этой строки прописан в условии задачи
            } else {
                unite(j, k); // В противном случае объединяем множество, содержащее i-ый элемент, и множество, содержащее j-ый.
            }
        }
        if (query == "CHECK") { // если запрос - это "CHECK", то нужно проверить 2 элемента, на принадлежность одному множеству
            int j, k;
            cin >> j >> k; // считываем j и k
            if (get_p(j) == get_p(k)) { // Если номера множеств, к которым принадлежат j и k совпадают, то они уже лежат в одном множестве
                cout << "YES\n"; // Выводим "YES"
            } else {
                cout << "NO\n"; // В противном случае выводим "NO"
            }
        }
    }
}

int main() {
    solve();
    return 0;
}