Cod sursa(job #3264365)

Utilizator alexsimedreaAlexandru Simedrea alexsimedrea Data 20 decembrie 2024 19:10:09
Problema Arbore partial de cost minim Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.71 kb
#include <algorithm>
#include <fstream>
#include <vector>

std::vector<std::vector<std::pair<int, int>>> adj;
std::vector<int> key;
std::vector<int> parent;
std::vector<int> heap;
std::vector<int> heapPos;
int heapSize = 0;

void introduce_in_mst(int x) {
    for (const auto &edge: adj[x]) {
        int node = edge.first;
        int cost = edge.second;
        if (cost < key[node]) {
            key[node] = cost;
            parent[node] = x;
        }
    }
}

void push_down(int pos) {
    while (true) {
        int smallest = pos;
        int left = 2 * pos;
        int right = 2 * pos + 1;

        if (left <= heapSize && key[heap[left]] < key[heap[smallest]]) {
            smallest = left;
        }
        if (right <= heapSize && key[heap[right]] < key[heap[smallest]]) {
            smallest = right;
        }

        if (smallest == pos) {
            break;
        }

        std::swap(heap[pos], heap[smallest]);
        std::swap(heapPos[heap[pos]], heapPos[heap[smallest]]);
        pos = smallest;
    }
}

void bubble_up(int pos) {
    while (pos > 1 && key[heap[pos]] < key[heap[pos / 2]]) {
        std::swap(heap[pos], heap[pos / 2]);
        std::swap(heapPos[heap[pos]], heapPos[heap[pos / 2]]);
        pos = pos / 2;
    }
}

void insert_to_heap(int node) {
    heap[++heapSize] = node;
    heapPos[node] = heapSize;
    bubble_up(heapSize);
}

int extract_min() {
    int min_node = heap[1];
    std::swap(heap[1], heap[heapSize]);
    std::swap(heapPos[heap[1]], heapPos[heap[heapSize]]);
    heapSize--;
    push_down(1);
    heapPos[min_node] = 0;
    return min_node;
}

int main() {
    std::ifstream fin("apm.in");
    std::ofstream fout("apm.out");

    int n, m;
    fin >> n >> m;

    adj.resize(n + 1);
    key.assign(n + 1, 1 << 30);
    parent.resize(n + 1);
    heap.resize(n + 1);
    heapPos.resize(n + 1);

    for (int i = 0; i < m; ++i) {
        int x, y, cost;
        fin >> x >> y >> cost;
        adj[x].emplace_back(y, cost);
        adj[y].emplace_back(x, cost);
    }

    key[1] = 0;
    introduce_in_mst(1);

    for (int i = 2; i <= n; ++i) {
        insert_to_heap(i);
    }

    int total_cost = 0;
    std::vector<std::pair<int, int>> mst_edges;

    for (int i = 1; i < n; ++i) {
        int u = extract_min();
        introduce_in_mst(u);
        total_cost += key[u];
        mst_edges.emplace_back(u, parent[u]);

        for (const auto &edge: adj[u]) {
            int v = edge.first;
            if (heapPos[v]) {
                bubble_up(heapPos[v]);
            }
        }
    }

    fout << total_cost << '\n';
    fout << n - 1 << '\n';
    for (const auto &edge: mst_edges) {
        fout << edge.first << " " << edge.second << '\n';
    }

    fin.close();
    fout.close();
    return 0;
}