Grafo:              Arestas ordenadas:
  0---1  peso 4     (0,3)=1, (2,3)=2, (0,1)=4,
  |\ /|             (1,3)=5, (1,2)=6
  1  5 6
  | /\ |            Passo 1: (0,3)=1 → 0 e 3 em componentes diferentes → ACEITA
  3---2  peso 2     Passo 2: (2,3)=2 → 2 e 3 em componentes diferentes → ACEITA
  peso 1            Passo 3: (0,1)=4 → 0 e 1 em componentes diferentes → ACEITA
                    → MST completa (3 arestas para 4 vértices) ✓
MST:
  0   1             Peso total: 1 + 2 + 4 = 7
  |   |
  3---2

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

/// Estrutura Union-Find com compressão de caminho e união por rank.
const UnionFind = struct {
    pai: []u32,
    rank: []u32,
    allocator: Allocator,

    pub fn init(allocator: Allocator, n: u32) !UnionFind {
        const pai = try allocator.alloc(u32, n);
        const rank = try allocator.alloc(u32, n);
        for (0..n) |i| {
            pai[i] = @intCast(i);
            rank[i] = 0;
        }
        return .{ .pai = pai, .rank = rank, .allocator = allocator };
    }

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

    /// Encontra o representante do conjunto (com compressão de caminho).
    pub fn find(self: *UnionFind, x: u32) u32 {
        if (self.pai[x] != x) {
            self.pai[x] = self.find(self.pai[x]);
        }
        return self.pai[x];
    }

    /// Une dois conjuntos (por rank).
    pub fn unir(self: *UnionFind, x: u32, y: u32) bool {
        const rx = self.find(x);
        const ry = self.find(y);
        if (rx == ry) return false; // já no mesmo conjunto

        if (self.rank[rx] < self.rank[ry]) {
            self.pai[rx] = ry;
        } else if (self.rank[rx] > self.rank[ry]) {
            self.pai[ry] = rx;
        } else {
            self.pai[ry] = rx;
            self.rank[rx] += 1;
        }
        return true;
    }
};

const Aresta = struct {
    u: u32,
    v: u32,
    peso: i64,
};

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

/// Algoritmo de Kruskal — retorna as arestas da MST e o peso total.
pub fn kruskal(
    allocator: Allocator,
    num_vertices: u32,
    arestas: []Aresta,
) !struct { mst: []Aresta, peso_total: i64 } {
    // Ordena arestas por peso
    std.mem.sort(Aresta, arestas, {}, compararArestas);

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

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

    for (arestas) |aresta| {
        // Se não forma ciclo, adiciona à MST
        if (uf.unir(aresta.u, aresta.v)) {
            try mst.append(aresta);
            peso_total += aresta.peso;

            // MST completa quando tem V-1 arestas
            if (mst.items.len == num_vertices - 1) break;
        }
    }

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

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

    var arestas = [_]Aresta{
        .{ .u = 0, .v = 1, .peso = 4 },
        .{ .u = 0, .v = 3, .peso = 1 },
        .{ .u = 1, .v = 2, .peso = 6 },
        .{ .u = 1, .v = 3, .peso = 5 },
        .{ .u = 2, .v = 3, .peso = 2 },
    };

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

    try stdout.print("Árvore Geradora Mínima (Kruskal):\n", .{});
    for (resultado.mst) |a| {
        try stdout.print("  {d} --- {d}  (peso: {d})\n", .{ a.u, a.v, a.peso });
    }
    try stdout.print("Peso total: {d}\n", .{resultado.peso_total});
}