// Árboles de expansión mínima: Kruskal, Prim, agrupamiento y las ideas que fallan.
#pragma once

#include <algorithm>
#include <cstddef>
#include <functional>
#include <limits>
#include <numeric>
#include <queue>
#include <tuple>
#include <utility>
#include <vector>

struct Edge {
    int u;
    int v;
    int weight;
};

inline long long total_weight(const std::vector<Edge>& edges) {
    long long total = 0;
    for (const Edge& e : edges) total += e.weight;
    return total;
}

class UnionFind {
public:
    explicit UnionFind(int n) : parent_(static_cast<std::size_t>(n)), size_(static_cast<std::size_t>(n), 1), sets_(n) {
        std::iota(parent_.begin(), parent_.end(), 0);
    }

    int find(int x) {
        while (parent_[x] != x) {
            parent_[x] = parent_[parent_[x]];  // compresión por mitades
            x = parent_[x];
        }
        return x;
    }

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

    int sets() const { return sets_; }

private:
    std::vector<int> parent_;
    std::vector<int> size_;
    int sets_;
};

// Ordena por peso. Con pesos iguales respeta el orden de la lista.
inline std::vector<Edge> sorted_by_weight(std::vector<Edge> edges) {
    std::stable_sort(edges.begin(), edges.end(),
                     [](const Edge& a, const Edge& b) { return a.weight < b.weight; });
    return edges;
}

// Devuelve un bosque de expansión mínima: un árbol por cada componente del grafo.
// Si el grafo es conexo, tiene n - 1 aristas.
inline std::vector<Edge> kruskal(int n, const std::vector<Edge>& edges) {
    UnionFind sets(n);
    std::vector<Edge> forest;
    for (const Edge& e : sorted_by_weight(edges)) {
        // Invariante: forest está contenido en algún árbol de expansión mínima.
        if (sets.unite(e.u, e.v)) forest.push_back(e);
    }
    return forest;
}

struct Neighbor {
    int to;
    int weight;
};

inline std::vector<std::vector<Neighbor>> adjacency(int n, const std::vector<Edge>& edges) {
    std::vector<std::vector<Neighbor>> adj(static_cast<std::size_t>(n));
    for (const Edge& e : edges) {
        adj[e.u].push_back({e.v, e.weight});
        adj[e.v].push_back({e.u, e.weight});
    }
    return adj;
}

struct PrimResult {
    std::vector<Edge> tree;
    long long pushes = 0;      // entradas que entraron a la cola
    long long stale_pops = 0;  // entradas que salieron y se descartaron
};

// Crece desde source. std::priority_queue no puede bajar la clave de una entrada, así que se
// usa borrado perezoso: las entradas viejas se quedan y se descartan cuando salen.
inline PrimResult prim(int n, const std::vector<Edge>& edges, int source) {
    using Entry = std::tuple<int, int, int>;  // peso, nodo de afuera, nodo del árbol
    auto adj = adjacency(n, edges);
    std::vector<bool> in_tree(static_cast<std::size_t>(n), false);
    std::priority_queue<Entry, std::vector<Entry>, std::greater<Entry>> heap;
    PrimResult result;

    auto add_node = [&](int u) {
        in_tree[u] = true;
        for (const Neighbor& nb : adj[u]) {
            if (!in_tree[nb.to]) {
                heap.push({nb.weight, nb.to, u});
                result.pushes++;
            }
        }
    };

    add_node(source);
    while (!heap.empty()) {
        auto [weight, v, from] = heap.top();
        heap.pop();
        if (in_tree[v]) {
            result.stale_pops++;  // sus dos extremos ya están adentro: no cruza el corte
            continue;
        }
        result.tree.push_back({from, v, weight});
        add_node(v);
    }
    return result;
}

// Distancias desde source con Dijkstra. También devuelve la arista con que se fijó cada nodo.
struct DijkstraResult {
    std::vector<long long> dist;
    std::vector<Edge> tree;
};

inline DijkstraResult dijkstra(int n, const std::vector<Edge>& edges, int source) {
    using Entry = std::tuple<long long, int, int, int>;  // distancia, nodo de afuera, nodo del árbol, peso
    auto adj = adjacency(n, edges);
    std::vector<bool> in_tree(static_cast<std::size_t>(n), false);
    std::priority_queue<Entry, std::vector<Entry>, std::greater<Entry>> heap;
    DijkstraResult result;
    result.dist.assign(static_cast<std::size_t>(n), -1);

    auto add_node = [&](int u, long long du) {
        in_tree[u] = true;
        result.dist[u] = du;
        for (const Neighbor& nb : adj[u]) {
            if (!in_tree[nb.to]) heap.push({du + nb.weight, nb.to, u, nb.weight});
        }
    };

    add_node(source, 0);
    while (!heap.empty()) {
        auto [d, v, from, weight] = heap.top();
        heap.pop();
        if (in_tree[v]) continue;
        result.tree.push_back({from, v, weight});
        add_node(v, d);
    }
    return result;
}

// Idea que falla: las n - 1 aristas más baratas, sin mirar si forman ciclos.
inline std::vector<Edge> cheapest_edges_only(int n, const std::vector<Edge>& edges) {
    std::vector<Edge> sorted = sorted_by_weight(edges);
    std::size_t count = std::min(sorted.size(), static_cast<std::size_t>(n - 1));
    sorted.resize(count);
    return sorted;
}

// Idea que falla: cada nodo se queda con su arista más barata.
inline std::vector<Edge> cheapest_per_node(int n, const std::vector<Edge>& edges) {
    std::vector<int> best(static_cast<std::size_t>(n), -1);
    for (int i = 0; i < static_cast<int>(edges.size()); i++) {
        for (int end : {edges[i].u, edges[i].v}) {
            if (best[end] == -1 || edges[i].weight < edges[best[end]].weight) best[end] = i;
        }
    }
    std::vector<int> chosen;
    for (int i : best) {
        if (i != -1) chosen.push_back(i);
    }
    std::sort(chosen.begin(), chosen.end());
    chosen.erase(std::unique(chosen.begin(), chosen.end()), chosen.end());
    std::vector<Edge> result;
    for (int i : chosen) result.push_back(edges[i]);
    return result;
}

inline int count_components(int n, const std::vector<Edge>& edges) {
    UnionFind sets(n);
    for (const Edge& e : edges) sets.unite(e.u, e.v);
    return sets.sets();
}

// Un árbol de expansión de n nodos: n - 1 aristas que conectan todo.
inline bool is_spanning_tree(int n, const std::vector<Edge>& edges) {
    return static_cast<int>(edges.size()) == n - 1 && count_components(n, edges) == 1;
}

// Fuerza bruta: todos los árboles de expansión. Solo para grafos pequeños.
inline std::vector<std::vector<int>> all_spanning_trees(int n, const std::vector<Edge>& edges) {
    std::vector<std::vector<int>> trees;
    std::vector<int> chosen;
    std::function<void(int)> go = [&](int i) {
        if (static_cast<int>(chosen.size()) == n - 1) {
            std::vector<Edge> picked;
            for (int c : chosen) picked.push_back(edges[c]);
            if (count_components(n, picked) == 1) trees.push_back(chosen);
            return;
        }
        if (i == static_cast<int>(edges.size())) return;
        chosen.push_back(i);
        go(i + 1);
        chosen.pop_back();
        go(i + 1);
    };
    go(0);
    return trees;
}

// Fuerza bruta: los árboles de expansión de costo mínimo, como listas de índices de aristas.
inline std::vector<std::vector<int>> all_minimum_trees(int n, const std::vector<Edge>& edges) {
    std::vector<std::vector<int>> best;
    long long best_cost = std::numeric_limits<long long>::max();
    for (const std::vector<int>& tree : all_spanning_trees(n, edges)) {
        long long cost = 0;
        for (int c : tree) cost += edges[c].weight;
        if (cost < best_cost) {
            best_cost = cost;
            best.clear();
        }
        if (cost == best_cost) best.push_back(tree);
    }
    return best;
}

// Grupos de una lista de aristas: group[v] es el número del grupo de v, por orden de aparición.
inline std::vector<int> groups_of(int n, const std::vector<Edge>& edges) {
    UnionFind sets(n);
    for (const Edge& e : edges) sets.unite(e.u, e.v);
    std::vector<int> label(static_cast<std::size_t>(n), -1);
    std::vector<int> group(static_cast<std::size_t>(n));
    int next = 0;
    for (int v = 0; v < n; v++) {
        int root = sets.find(v);
        if (label[root] == -1) label[root] = next++;
        group[v] = label[root];
    }
    return group;
}

// Agrupar en k grupos: quitar las k - 1 aristas más caras de un árbol de expansión mínima.
inline std::vector<int> max_spacing_groups(int n, const std::vector<Edge>& edges, int k) {
    std::vector<Edge> kept = sorted_by_weight(kruskal(n, edges));
    kept.resize(kept.size() - static_cast<std::size_t>(k - 1));
    return groups_of(n, kept);
}

// Separación: la arista más barata entre dos grupos distintos. -1 si ninguna los une.
inline int spacing(const std::vector<Edge>& edges, const std::vector<int>& group) {
    int best = -1;
    for (const Edge& e : edges) {
        if (group[e.u] != group[e.v] && (best == -1 || e.weight < best)) best = e.weight;
    }
    return best;
}

// Fuerza bruta: la mayor separación entre todas las particiones en k grupos no vacíos.
// partitions, si no es nulo, recibe cuántas particiones se probaron.
inline int best_spacing_brute(int n, const std::vector<Edge>& edges, int k, long long* partitions = nullptr) {
    std::vector<int> group(static_cast<std::size_t>(n), 0);
    int best = -1;
    long long seen = 0;
    // Cada partición se genera una vez: el grupo de v es a lo sumo uno más que los ya usados.
    std::function<void(int, int)> go = [&](int v, int used) {
        if (v == n) {
            if (used == k) {
                seen++;
                best = std::max(best, spacing(edges, group));
            }
            return;
        }
        for (int g = 0; g <= used && g < k; g++) {
            group[v] = g;
            go(v + 1, std::max(used, g + 1));
        }
    };
    go(0, 0);
    if (partitions != nullptr) *partitions = seen;
    return best;
}