Cod sursa(job #2681427)

Utilizator AlexVulpoiuAlexandru Vulpoiu AlexVulpoiu Data 5 decembrie 2020 14:12:00
Problema Arbore partial de cost minim Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.74 kb
/// Prim, O(M * log(N))
#include <fstream>
#include <vector>

using namespace std;

struct Edge {
    int x, y, cost;
};

int n, m, i, j, k, cmin;
bool viz[200005];
Edge mc, aux, heap[200005];
vector<pair<int, int>> arbore, v[200005];

void pushup(int p) {
    int val;

    val = heap[p].cost;
    while(p > 1 && val < heap[p / 2].cost) {
        swap(heap[p], heap[p / 2]);
        p /= 2;
    }
}

void pushdown(int p) {
    int fiu;

    fiu = 1;
    while(fiu) {
        fiu = 0;
        if(p * 2 <= m) {
            fiu = p * 2;
            if(p * 2 + 1 <= m && heap[p * 2 + 1].cost < heap[p * 2].cost)
                fiu = p * 2 + 1;
            if(heap[p].cost < heap[fiu].cost)
                fiu = 0;
        }
        if(fiu) {
            swap(heap[p], heap[fiu]);
            p = fiu;
        }
    }
}

int main()
{
    ifstream f("apm.in");
    ofstream g("apm.out");

    f >> n >> m;
    while(m) {
        f >> i >> j >> k;
        v[i].emplace_back(j, k);    /// muchie i->j, cu costul k
        v[j].emplace_back(i, k);    /// muchie j->i, cu costul k
        m--;
    }
    f.close();

    for(i = 1; i <= n; i++)
        viz[i] = false;     /// initial, niciun nod nu este vizitat
    viz[1] = true;          /// incepem constructia arborelui din nodul 1
    m = 0;                  /// dimensiune heap
    for(i = 0; i < v[1].size(); i++) {
        mc.x = 1;
        mc.y = v[1][i].first;
        mc.cost = v[1][i].second;
        heap[++m] = mc;     /// adaug muchiile incidente cu 1 in heap
        pushup(m);
    }
    cmin = 0;   /// cost arbore
    for(k = 1; k < n; k++) {
        mc = heap[1];                       /// extrag din heap muchia de cost minim, care are o extremitate in nodurile deja vizitate
        cmin += mc.cost;                    /// cresc costul arborelui
        heap[1] = heap[m--];
        pushdown(1);                    /// elimin din heap muchia extrasa
        arbore.emplace_back(mc.x, mc.y);    /// adaug muchia in arbore
        viz[mc.y] = true;                   /// vizitez cealalta extremitate
        for(i = 0; i < v[mc.y].size(); i++)
            if(!viz[v[mc.y][i].first]) {    /// adaug in heap muchiile care au o extremitate in y, iar cealalta nu este vizitata
                aux.x = mc.y;
                aux.y = v[mc.y][i].first;
                aux.cost = v[mc.y][i].second;
                heap[++m] = aux;
                pushup(m);
            }
        while(viz[heap[1].x] && viz[heap[1].y]) {   /// cat timp prima muchia din heap are ambele varfuri vizitate, o elimin
            heap[1] = heap[m--];
            pushdown(1);
        }
    }

    g << cmin << '\n' << n - 1 << '\n';
    for(i = 0; i < arbore.size(); i++)
        g << arbore[i].first << ' ' << arbore[i].second << '\n';
    g.close();

    return 0;
}