Cod sursa(job #3130675)

Utilizator mateicalin2002Ceausu Matei Calin mateicalin2002 Data 18 mai 2023 13:41:35
Problema Arbore partial de cost minim Scor 100
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.58 kb
#include <bits/stdc++.h>

using namespace std;

#define NMAX 200005

class DisjointSet {
private:
    vector<int> parent;
    vector<int> size;

public:
    DisjointSet(int nodes)
        : parent(nodes + 1)
        , size(nodes + 1) {
        for (int node = 1; node <= nodes; ++node) {
            parent[node] = node;
            size[node] = 1;
        }
    }

    int setOf(int node) {
        if (node == parent[node]) {
            return node;
        }

        parent[node] = setOf(parent[node]);
        return parent[node];
    }

    void _union(int x, int y) {
        int rx = setOf(x), ry = setOf(y);

        if (size[rx] <= size[ry]) {
            size[ry] += size[rx];
            parent[rx] = ry;
        } else {
            size[rx] += size[ry];
            parent[ry] = rx;
        }
    }
};

struct Edge {
    int node;
    int neigh;
    int w;

    Edge() { }
    Edge(int node, int neigh, int w)
        : node(node)
        , neigh(neigh)
        , w(w) { }
};

struct MSTResult {
    int cost;
    vector<pair<int, int>> edges;

    MSTResult(int cost, const vector<pair<int, int>>& edges)
        : cost(cost)
        , edges(edges) { }
};

class Task {
public:
    void solve() {
        read_input();
        print_output(get_result());
    }

private:
    int n, m;

    vector<Edge> edges;

    void read_input() {
        ifstream fin("apm.in");
        fin >> n >> m;
        for (int i = 1, x, y, w; i <= m; i++) {
            fin >> x >> y >> w;
            edges.push_back(Edge{x, y, w});
        }
        fin.close();
    }

    MSTResult get_result() {
        sort(edges.begin(), edges.end(), [](const Edge& a, const Edge& b) {
            return a.w < b.w;
        });

        DisjointSet disjointset(n + 5);

        int cost = 0, cnt = 0;
        vector<pair<int, int>> mst;

        for (int i = 0; i < edges.size() && cnt < n - 1; ++i) {
            int x = edges[i].node;
            int y = edges[i].neigh;
            int w = edges[i].w;

            if (disjointset.setOf(x) != disjointset.setOf(y)) {
                disjointset._union(x, y);
                cost += w;
                mst.push_back({x, y});
                ++cnt;
            }
        }

        return {cost, mst};
    }

    void print_output(const MSTResult& res) {
        ofstream fout("apm.out");
        fout << res.cost << "\n";
        fout << res.edges.size() << "\n";
        for (const auto& [x, y] : res.edges) {
            fout << x << " " << y << "\n";
        }
        fout.close();
    }
};

int main() {
    auto* task = new (nothrow) Task();
    task->solve();
    delete task;
    return 0;
}