Inicialmente: cada elemento é seu próprio conjunto
  {0} {1} {2} {3} {4} {5} {6} {7}

Union(0, 1):  {0, 1} {2} {3} {4} {5} {6} {7}
Union(2, 3):  {0, 1} {2, 3} {4} {5} {6} {7}
Union(0, 3):  {0, 1, 2, 3} {4} {5} {6} {7}
Union(5, 6):  {0, 1, 2, 3} {4} {5, 6} {7}

Find(1) == Find(3) → true  (mesmo conjunto)
Find(1) == Find(5) → false (conjuntos diferentes)

Árvore interna (com compressão):
    0           5
   / \         |
  1   2        6
      |
      3

Após Find(3) com compressão:
    0           5
   /|\         |
  1  2  3      6

const std = @import("std");
const Allocator = std.mem.Allocator;

/// Union-Find com compressão de caminho e união por rank.
pub const UnionFind = struct {
    const Self = @This();

    pai: []u32,
    rank: []u8,
    tamanho: []u32,
    num_conjuntos: u32,
    allocator: Allocator,

    pub fn init(allocator: Allocator, n: u32) !Self {
        var pai = try allocator.alloc(u32, n);
        var rank_arr = try allocator.alloc(u8, n);
        var tam = try allocator.alloc(u32, n);

        for (0..n) |i| {
            pai[i] = @intCast(i);   // Cada elemento é raiz de si mesmo
            rank_arr[i] = 0;
            tam[i] = 1;
        }

        return .{
            .pai = pai,
            .rank = rank_arr,
            .tamanho = tam,
            .num_conjuntos = n,
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *Self) void {
        self.allocator.free(self.pai);
        self.allocator.free(self.rank);
        self.allocator.free(self.tamanho);
    }

    /// Find com compressão de caminho — quase O(1).
    pub fn find(self: *Self, x: u32) u32 {
        if (self.pai[x] != x) {
            self.pai[x] = self.find(self.pai[x]); // Compressão
        }
        return self.pai[x];
    }

    /// Union por rank — quase O(1).
    /// Retorna true se os conjuntos foram unidos (eram diferentes).
    pub fn unir(self: *Self, a: u32, b: u32) bool {
        const raiz_a = self.find(a);
        const raiz_b = self.find(b);

        if (raiz_a == raiz_b) return false; // Já no mesmo conjunto

        // União por rank: árvore menor fica sob a maior
        if (self.rank[raiz_a] < self.rank[raiz_b]) {
            self.pai[raiz_a] = raiz_b;
            self.tamanho[raiz_b] += self.tamanho[raiz_a];
        } else if (self.rank[raiz_a] > self.rank[raiz_b]) {
            self.pai[raiz_b] = raiz_a;
            self.tamanho[raiz_a] += self.tamanho[raiz_b];
        } else {
            self.pai[raiz_b] = raiz_a;
            self.tamanho[raiz_a] += self.tamanho[raiz_b];
            self.rank[raiz_a] += 1;
        }

        self.num_conjuntos -= 1;
        return true;
    }

    /// Verifica se dois elementos estão no mesmo conjunto.
    pub fn conectados(self: *Self, a: u32, b: u32) bool {
        return self.find(a) == self.find(b);
    }

    /// Retorna o tamanho do conjunto que contém x.
    pub fn tamanhoConjunto(self: *Self, x: u32) u32 {
        const raiz = self.find(x);
        return self.tamanho[raiz];
    }
};

// Aplicação: Kruskal para Árvore Geradora Mínima (MST)
const Aresta = struct {
    u: u32,
    v: u32,
    peso: i64,
};

fn compararPeso(_: void, a: Aresta, b: Aresta) bool {
    return a.peso < b.peso;
}

fn kruskal(arestas: []Aresta, num_vertices: u32, allocator: Allocator) !struct {
    mst: []Aresta,
    custo_total: i64,
} {
    // Ordena arestas por peso
    std.mem.sort(Aresta, arestas, {}, compararPeso);

    var uf = try UnionFind.init(allocator, num_vertices);
    defer uf.deinit();

    var mst = std.ArrayList(Aresta).init(allocator);
    var custo: i64 = 0;

    for (arestas) |aresta| {
        if (uf.unir(aresta.u, aresta.v)) {
            try mst.append(aresta);
            custo += aresta.peso;
            if (mst.items.len == num_vertices - 1) break;
        }
    }

    return .{
        .mst = try mst.toOwnedSlice(),
        .custo_total = custo,
    };
}

pub fn main() !void {
    const stdout = std.io.getStdOut().writer();
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    // Union-Find básico
    try stdout.writeAll("=== Union-Find Básico ===\n");
    var uf = try UnionFind.init(allocator, 8);
    defer uf.deinit();

    _ = uf.unir(0, 1);
    _ = uf.unir(2, 3);
    _ = uf.unir(0, 3);
    _ = uf.unir(5, 6);

    try stdout.print("0 e 3 conectados: {}\n", .{uf.conectados(0, 3)});
    try stdout.print("1 e 2 conectados: {}\n", .{uf.conectados(1, 2)});
    try stdout.print("0 e 5 conectados: {}\n", .{uf.conectados(0, 5)});
    try stdout.print("Conjuntos: {d}\n", .{uf.num_conjuntos});
    try stdout.print("Tamanho do conjunto de 0: {d}\n", .{uf.tamanhoConjunto(0)});

    // Kruskal MST
    try stdout.writeAll("\n=== Kruskal — Árvore Geradora Mínima ===\n");
    var arestas = [_]Aresta{
        .{ .u = 0, .v = 1, .peso = 4 },
        .{ .u = 0, .v = 3, .peso = 8 },
        .{ .u = 1, .v = 2, .peso = 2 },
        .{ .u = 1, .v = 4, .peso = 3 },
        .{ .u = 2, .v = 5, .peso = 7 },
        .{ .u = 3, .v = 4, .peso = 1 },
        .{ .u = 4, .v = 5, .peso = 6 },
    };

    const resultado = try kruskal(&arestas, 6, allocator);
    defer allocator.free(resultado.mst);

    try stdout.writeAll("Arestas na MST:\n");
    for (resultado.mst) |a| {
        try stdout.print("  {d} — {d} (peso {d})\n", .{ a.u, a.v, a.peso });
    }
    try stdout.print("Custo total: {d}\n", .{resultado.custo_total});
}