Cod sursa(job #2651167)

Utilizator marius004scarlat marius marius004 Data 21 septembrie 2020 18:00:22
Problema Lazy Scor 50
Compilator cpp-64 Status done
Runda Arhiva de probleme Marime 1.89 kb
#include <fstream>
#include <vector>
#include <algorithm>

using namespace std;

ifstream f("lazy.in");
ofstream g("lazy.out");

struct Edge {
    int x, y, c1, c2, index;

    friend istream& operator >> (istream& stream, Edge& e) {
        stream >> e.x >> e.y >> e.c1 >> e.c2;
        return stream;
    }

    friend ostream& operator << (ostream& stream, Edge& e) {
        stream << e.x << ' ' << e.y << ' ' << e.c1 << ' ' << e.c2;
        return stream;
    }

    friend bool operator < (const Edge& a, const Edge& b) {
        return (a.c1 != b.c1 ? (a.c1 < b.c1) : (a.c2 > b.c2));
    }
};

template <int SIZE>
class DisjointSet {

private:

    vector < int > tree;

    void collapse(int node, const int& root) {

        while(node != root) {
            int aux = tree[node];

            tree[node] = root;
            node = aux;
        }

    }

public:

    DisjointSet() {
        tree.resize(SIZE, -1);
    }

    int Find(const int& node) {

        int root = node;

        while(tree[root] > 0)
            root = tree[root];

        collapse(node, root);

        return root;
    }

    void Union(int X, int Y) {

        if(-tree[X] < -tree[Y])
            swap(X, Y);

        tree[X] += tree[Y];
        tree[Y] = X;
    }
};

const int NMAX = 200'001;

int N, M;
vector < Edge > edges;

void kruskal() {

    DisjointSet < NMAX > disjoint;

    sort(edges.begin(), edges.end());

    int cnt = 0;
    for(const Edge& e : edges) {

        const int rootX = disjoint.Find(e.x);
        const int rootY = disjoint.Find(e.y);

        if(rootX != rootY) {

            disjoint.Union(rootX, rootY);

            g << e.index << "\n";
            cnt++;

            if(cnt == N - 1)
                break;
        }
    }
}

int main() {

    f >> N >> M;

    edges.resize(M);

    for(int i = 0;i < M;++i) {
        f >> edges[i];
        edges[i].index = i + 1;
    }

    kruskal();

    return 0;
}