Graph/tests/yosupo_mst.test.cpp

View this file on GitHub
Last update: 2024-03-17 02:51:32+00:00
Problem: https://judge.yosupo.jp/problem/minimum_spanning_tree

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/minimum_spanning_tree"

#include "../../template.h"
#include "../mst.h"

struct E : Edge {
    int id;
};

void solve() {
    int n, m; cin >> n >> m;
    vector<E> edges(m);
    REP(i,m) {
        auto& e = edges[i];
        cin >> e.u >> e.v >> e.c;
        e.id = i;
    }
    auto g = mst<E>(n, edges);
    cout << g.first << '\n';
    for (auto& e : g.second) cout << e.id << ' ';
    cout << '\n';
}

#line 1 "Graph/tests/yosupo_mst.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/minimum_spanning_tree"

#line 1 "template.h"
#include <bits/stdc++.h>
using namespace std;

#define FOR(i,a,b) for(int i=(a),_b=(b); i<=_b; i++)
#define FORD(i,a,b) for(int i=(a),_b=(b); i>=_b; i--)
#define REP(i,a) for(int i=0,_a=(a); i<_a; i++)
#define EACH(it,a) for(__typeof(a.begin()) it = a.begin(); it != a.end(); ++it)

#define DEBUG(x) { cout << #x << " = "; cout << (x) << endl; }
#define PR(a,n) { cout << #a << " = "; FOR(_,1,n) cout << a[_] << ' '; cout << endl; }
#define PR0(a,n) { cout << #a << " = "; REP(_,n) cout << a[_] << ' '; cout << endl; }

#define sqr(x) ((x) * (x))

// For printing pair, container, etc.
// Copied from https://quangloc99.github.io/2021/07/30/my-CP-debugging-template.html
template<class U, class V> ostream& operator << (ostream& out, const pair<U, V>& p) {
    return out << '(' << p.first << ", " << p.second << ')';
}

template<class Con, class = decltype(begin(declval<Con>()))>
typename enable_if<!is_same<Con, string>::value, ostream&>::type
operator << (ostream& out, const Con& con) {
    out << '{';
    for (auto beg = con.begin(), it = beg; it != con.end(); it++) {
        out << (it == beg ? "" : ", ") << *it;
    }
    return out << '}';
}
template<size_t i, class T> ostream& print_tuple_utils(ostream& out, const T& tup) {
    if constexpr(i == tuple_size<T>::value) return out << ")"; 
    else return print_tuple_utils<i + 1, T>(out << (i ? ", " : "(") << get<i>(tup), tup); 
}
template<class ...U> ostream& operator << (ostream& out, const tuple<U...>& t) {
    return print_tuple_utils<0, tuple<U...>>(out, t);
}

mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
long long get_rand(long long r) {
    return uniform_int_distribution<long long> (0, r-1)(rng);
}

template<typename T>
vector<T> read_vector(int n) {
    vector<T> res(n);
    for (int& x : res) cin >> x;
    return res;
}

void solve();

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    solve();
    return 0;
}
#line 1 "Graph/mst.h"
// MST. 0-based index
//
// Returns:
// {mst cost, edges in mst}
//
// If graph is not connected, returns forest (number of edges will be < n-1)

#line 1 "DataStructure/DSU/DisjointSet.h"
// DisjointSet {{{
struct DSU {
    vector<int> lab;

    DSU(int n) : lab(n+1, -1) {}

    int getRoot(int u) {
        if (lab[u] < 0) return u;
        return lab[u] = getRoot(lab[u]);
    }

    bool merge(int u, int v) {
        u = getRoot(u); v = getRoot(v);
        if (u == v) return false;
        if (lab[u] > lab[v]) swap(u, v);
        lab[u] += lab[v];
        lab[v] = u;
        return true;
    }

    bool same_component(int u, int v) {
        return getRoot(u) == getRoot(v);
    }

    int component_size(int u) {
        return -lab[getRoot(u)];
    }
};
// }}}
#line 9 "Graph/mst.h"

// MST {{{
using ll = long long;
template<typename EdgeT>
std::pair<ll, std::vector<EdgeT>> mst(
        int n,
        std::vector<EdgeT> edges) {
    std::sort(edges.begin(), edges.end());

    DSU dsu(n + 1);  // tolerate 1-based index
    ll total = 0;
    vector<EdgeT> tree;
    for (const auto& e : edges) {
        if (dsu.merge(e.u, e.v)) {
            total += e.c;
            tree.push_back(e);
        }
    }
    return {total, tree};
}
struct Edge {
    int u, v;
    ll c;
};
bool operator < (const Edge& a, const Edge& b) {
    return a.c < b.c;
}
ostream& operator << (ostream& out, const Edge& e) {
    out << e.u << " - " << e.v << " [" << e.c << ']';
    return out;
}
// }}}
#line 5 "Graph/tests/yosupo_mst.test.cpp"

struct E : Edge {
    int id;
};

void solve() {
    int n, m; cin >> n >> m;
    vector<E> edges(m);
    REP(i,m) {
        auto& e = edges[i];
        cin >> e.u >> e.v >> e.c;
        e.id = i;
    }
    auto g = mst<E>(n, edges);
    cout << g.first << '\n';
    for (auto& e : g.second) cout << e.id << ' ';
    cout << '\n';
}