Cod sursa(job #2205277)

Utilizator Andreiii500Andrei Puiu Andreiii500 Data 18 mai 2018 17:27:12
Problema Arbore partial de cost minim Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 3.39 kb
/// PRIM (N LOG M)
#include<cstdio>
#include<iostream>
#include<fstream>
#include<vector>
#include<queue>
#include<cassert>
using namespace std;

ifstream in("apm.in");
//ofstream out("apm.out");

const int INF = 1001;
const int N = 200001;
const int START_NODE = 1;
const int UNKNOWN_PARENT = -2;
const bool DEBUG = false;

struct node{
    int index;
    int cost_to_reach = INF;
    int parent_to_reach = UNKNOWN_PARENT;
    bool is_in_the_tree = false;
};

struct edge{
    int cost;
    int index;
    int index_parent;
};

auto cmp = [](edge a, edge b){return a.cost > b.cost;};

vector<pair<int, int> > neighbours[N]; /// index_neighbour, cost
priority_queue<edge, vector<edge>, decltype(cmp) > edges(cmp);
node v[N];
int n,m;

void expandTree(){
    edge best_edge;
    best_edge.index = -1;
    while(best_edge.index == -1 && !edges.empty()){
        edge top_edge = edges.top();
        edges.pop();

        if(v[top_edge.index].is_in_the_tree == false)
            best_edge = top_edge;
    }

    /// If we cannot join a new node, we shouldn't call this function :)
    assert(best_edge.index != -1);

    /// Add it into the tree
    //cout<<"Adding Node "<<best_edge.index<<" to the queue\n";
    v[best_edge.index].is_in_the_tree = true;
    v[best_edge.index].cost_to_reach =  best_edge.cost;
    //cout<<"cost to reach = "<<best_edge.cost<<"\n";
    if(best_edge.index != START_NODE)
        v[best_edge.index].parent_to_reach =  best_edge.index_parent;

    /// Check its neighbours and Update their cost_to_reach
    for(auto &neighbour : neighbours[best_edge.index])
        if(v[neighbour.first].is_in_the_tree == false)
        {
            edge this_edge;
            this_edge.cost = neighbour.second;
            this_edge.index = neighbour.first;
            this_edge.index_parent = best_edge.index;
            //cout<<"Adding ("<<this_edge.index_parent<<", "<<this_edge.index<<") "<<this_edge.cost<<"\n";

            edges.push(this_edge);
        }
}

int main()
{
    /// Input
    freopen("apm.out","w",stdout);

    in>>n>>m;
    for(int i=0; i<m; ++i){
        int index_1, index_2, cost;
        in>>index_1>>index_2>>cost;

        neighbours[index_1].push_back({index_2, cost});
        neighbours[index_2].push_back({index_1, cost});
    }

    /// Init
    for(int i=1; i<=n; ++i){
        v[i].index = i;
    }

    /// Solve
    edge first_edge;
    first_edge.cost = 0;
    first_edge.index = START_NODE;
    first_edge.index_parent = START_NODE;
    edges.push(first_edge);
    for(int i=1; i<=n; ++i){
        expandTree();

        if(DEBUG){
            cout<<"-------------------------------\n";
            for(int j=1; j<=n; ++j)
                if(v[j].is_in_the_tree)
                    cout<<"Node "<<j<<" is in the tree!\n";
        }
    }

    /// Print Solution
    int tree_cost = 0;
    for(int i=1; i<=n; ++i)
        tree_cost += v[i].cost_to_reach;
    //cout<<"Tree cost = ";
    cout<<tree_cost<<"\n";
    cout<<n-1<<"\n";
    for(int i=1; i<=n; ++i)
        if(v[i].parent_to_reach != UNKNOWN_PARENT)
            cout<<v[i].parent_to_reach<<" "<<i<<"\n";


    /// Test
    if(DEBUG){
        for(int i=1; i<=n; ++i){
            for(auto &neighbour : neighbours[i])
                cout<<i<<" "<<neighbour.first<<" "<<neighbour.second<<"\n";
        }
    }

    return 0;
}