// Bosque de union-find que cuenta los saltos por parent, con cada mejora opcional. Sirve para
// medir y dibujar. La clase del laboratorio tiene otra interfaz, y siempre usa las dos mejoras.
#ifndef COUNTING_FOREST_H
#define COUNTING_FOREST_H

#include <algorithm>
#include <cstdint>
#include <numeric>
#include <string>
#include <utility>
#include <vector>

// Invariantes: una raíz es su propio padre; si x no es raíz y se usa rango,
// rank[x] < rank[parent[x]]; sets_ es la cantidad de raíces.
class CountingForest {
public:
    CountingForest(int n, bool by_rank, bool compress)
        : parent_(n), rank_(n, 0), sets_(n), by_rank_(by_rank), compress_(compress) {
        std::iota(parent_.begin(), parent_.end(), 0);
    }

    int find(int x) {
        int root = x;
        while (parent_[root] != root) {
            root = parent_[root];
            hops_++;
        }
        while (compress_ && parent_[x] != root && x != root) {
            int next = parent_[x];
            parent_[x] = root;
            x = next;
        }
        return root;
    }

    bool unite(int a, int b) {
        int root_a = find(a);
        int root_b = find(b);
        if (root_a == root_b) {
            return false;
        }
        if (by_rank_) {
            if (rank_[root_a] < rank_[root_b]) {
                std::swap(root_a, root_b);
            }
            if (rank_[root_a] == rank_[root_b]) {
                rank_[root_a]++;
            }
            parent_[root_b] = root_a;
        } else {
            parent_[root_a] = root_b;
        }
        sets_--;
        return true;
    }

    int sets() const { return sets_; }
    std::int64_t hops() const { return hops_; }
    const std::vector<int>& parents() const { return parent_; }
    const std::vector<int>& ranks() const { return rank_; }

    // Profundidad de x, sin comprimir ni contar.
    int depth(int x) const {
        int steps = 0;
        while (parent_[x] != x) {
            x = parent_[x];
            steps++;
        }
        return steps;
    }

    int height() const {
        int best = 0;
        for (int x = 0; x < static_cast<int>(parent_.size()); x++) {
            best = std::max(best, depth(x));
        }
        return best;
    }

    // El árbol de root, escrito como en la semana 10: 0(1,2(3)).
    std::string shape(int root) const {
        std::string text = std::to_string(root);
        std::string children;
        for (int x = 0; x < static_cast<int>(parent_.size()); x++) {
            if (x != root && parent_[x] == root) {
                children += (children.empty() ? "" : ",") + shape(x);
            }
        }
        return children.empty() ? text : text + "(" + children + ")";
    }

    bool check_invariants() const {
        int roots = 0;
        for (int x = 0; x < static_cast<int>(parent_.size()); x++) {
            if (parent_[x] == x) {
                roots++;
            } else if (by_rank_ && rank_[x] >= rank_[parent_[x]]) {
                return false;
            }
        }
        return roots == sets_;
    }

private:
    std::vector<int> parent_;
    std::vector<int> rank_;
    int sets_;
    bool by_rank_;
    bool compress_;
    std::int64_t hops_ = 0;
};

#endif