Cod sursa(job #1976214)

Utilizator dani_aDaniel Alexandru dani_a Data 2 mai 2017 21:57:03
Problema Arbore partial de cost minim Scor 80
Compilator cpp Status done
Runda Arhiva educationala Marime 2.65 kb
#include <fstream>
#include <cstdint>
#include <map>
#include <memory>
#include <vector>
#include <set>

struct Input
{
    uint32_t nrOfNodes;
    uint32_t nrOfEdges;
    std::multimap<int16_t, std::pair<int32_t, int32_t>> edgedSortedByCost;
};

struct Output
{
    int32_t totalCost;
    std::vector<std::pair<int32_t, int32_t>> minimumSpanningTreeEdges;
};

Input readInput(std::ifstream& in)
{
    Input result;
    in >> result.nrOfNodes;
    in >> result.nrOfEdges;

    int32_t x, y; int16_t cost;

    for (int i = 0; i < result.nrOfEdges; ++i)
    {
        in >> x >> y >> cost;
        result.edgedSortedByCost.insert(std::make_pair(cost, std::make_pair(x, y)));
    }

    return result;
}

auto createSetVector(int32_t nrOfNodes)
{
    std::vector<std::shared_ptr<std::set<int32_t>>> result(nrOfNodes + 1);
    for (int i = 0; i <= nrOfNodes; ++i)
    {
        result[i] = std::make_shared<std::set<int32_t>>();
        result[i]->insert(i);
    }
    return result;
}

Output solveWithKruskal(Input& input)
{
    auto sets = createSetVector(input.nrOfNodes);

    auto nrOfOutputEdgesNeeded = input.nrOfNodes - 1;
    Output output;
    output.totalCost = 0;

    for (auto& costEdgesPair : input.edgedSortedByCost)
    {
        auto& cost = costEdgesPair.first;
        auto& edges = costEdgesPair.second;

        auto& firstSet = sets[edges.first];
        auto& secondSet = sets[edges.second];
        if (firstSet != secondSet)
        {
            output.totalCost += cost;
            output.minimumSpanningTreeEdges.push_back(edges);

            std::shared_ptr<std::set<int32_t>> smallerSet, biggerSet;

            if (firstSet->size() < secondSet->size())
            {
                smallerSet = firstSet;
                biggerSet = secondSet;
            }
            else
            {
                smallerSet = secondSet;
                biggerSet = firstSet;
            }

            for (auto& node : *smallerSet)
            {
                biggerSet->insert(node);
                sets[node] = biggerSet;
            }
        }
        if (output.minimumSpanningTreeEdges.size() >= nrOfOutputEdgesNeeded)
        {
            break;
        }
    }

    return output;
}

int main()
{
    std::ifstream in("apm.in");
    auto input = readInput(in);

    auto output = solveWithKruskal(input);

    std::ofstream out("apm.out");

    out << output.totalCost << '\n';
    out << output.minimumSpanningTreeEdges.size() << '\n';
    for (auto& pair : output.minimumSpanningTreeEdges)
    {
        out << pair.first << ' ' << pair.second << '\n';
    }
}