Graph/tests/directed_mst.test.cpp

View this file on GitHub
Last update: 2023-01-07 01:46:12+08:00
Problem: https://judge.yosupo.jp/problem/directedmst

Depends on

Code

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

#include <bits/stdc++.h>
using namespace std;

#include "../DirectedMST.h"

#define REP(i, a) for (int i = 0, _##i = (a); i < _##i; ++i)

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n, m, root; cin >> n >> m >> root;
    vector<Edge> edges(m);
    for (auto& e : edges) {
        cin >> e.u >> e.v >> e.cost;
    }
    auto [total, par] = directed_mst(n, root, edges);
    cout << total << endl;
    REP(i,n) cout << (i == root ? root : par[i]) << ' ';
    cout << endl;
    return 0;
}

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

#include <bits/stdc++.h>
using namespace std;

#line 1 "Graph/DirectedMST.h"
// include DSU_rollback.h

// Directed MST
// Index from 0
//
// Tested:
// - https://judge.yosupo.jp/problem/directedmst

#line 1 "DataStructure/DSU/DSU_rollback.h"
// Tested:
// - (dynamic connectivity) https://codeforces.com/gym/100551/problem/A
// - (used for directed MST) https://judge.yosupo.jp/problem/directedmst
//
// 0-based
// DSU with rollback {{{
struct Data {
    int time, u, par;  // before `time`, `par` = par[u]
};

struct DSU {
    vector<int> par;
    vector<Data> change;

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

    // find root of x.
    // if par[x] < 0 then x is a root, and its tree has -par[x] nodes
    int getRoot(int x) {
        while (par[x] >= 0)
            x = par[x];
        return x;
    }

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

    // join components containing x and y.
    // t should be current time. We use it to update `change`.
    bool join(int x, int y, int t) {
        x = getRoot(x);
        y = getRoot(y);
        if (x == y) return false;

        //union by rank
        if (par[x] < par[y]) swap(x, y); 
        //now x's tree has less nodes than y's tree
        change.push_back({t, y, par[y]});
        par[y] += par[x];
        change.push_back({t, x, par[x]});
        par[x] = y;
        return true;
    }

    // rollback all changes at time > t.
    void rollback(int t) {
        while (!change.empty() && change.back().time > t) {
            par[change.back().u] = change.back().par;
            change.pop_back();
        }
    }
};
// }}}
#line 10 "Graph/DirectedMST.h"
using ll = long long;
struct Edge {
    int u, v;  // directed, u -> v
    ll cost;
};
struct HeapNode {  // lazy skew heap node
    Edge key;
    HeapNode *l, *r;
    ll delta;

    void prop() {
        key.cost += delta;
        if (l) l->delta += delta;
        if (r) r->delta += delta;
        delta = 0;
    }
    Edge top() {
        prop();
        return key;
    }
};
HeapNode* merge(HeapNode *a, HeapNode *b) {
    if (!a) return b;
    if (!b) return a;
    a->prop(); b->prop();
    if (a->key.cost > b->key.cost) swap(a, b);
    swap(a->l, (a->r = merge(b, a->r)));
    return a;
}
void pop(HeapNode *&a) {
    a->prop();
    a = merge(a->l, a->r);
}

// return {cost, parent[i]}
// parent[root] = -1
// Not found -> return {-1, {}}
pair<ll, vector<int>> directed_mst(int n, int root, vector<Edge>& edges) {
    DSU dsu(n);
    int dsu_time = 0;
    vector<HeapNode*> heap(n);
    for (const Edge& e : edges) {
        heap[e.v] = merge(heap[e.v], new HeapNode{e});
    }

    ll res = 0;
    vector<int> seen(n, -1), path(n);
    seen[root] = root;
    vector<Edge> Q(n), in(n, {-1, -1});
    deque<tuple<int, int, vector<Edge>>> cycs;
    for (int s = 0; s < n; ++s) {
        int u = s, qi = 0, w;
        while (seen[u] < 0) {
            if (!heap[u]) return {-1, {}};
            Edge e = heap[u]->top();
            heap[u]->delta -= e.cost;
            pop(heap[u]);
            Q[qi] = e;
            path[qi++] = u;
            seen[u] = s;
            res += e.cost;
            u = dsu.getRoot(e.u);

            if (seen[u] == s) {
                HeapNode* cyc = 0;
                int end = qi;
                int time = dsu_time;
                do {
                    cyc = merge(cyc, heap[w = path[--qi]]);
                } while (dsu.join(u, w, ++dsu_time));

                u = dsu.getRoot(u);
                heap[u] = cyc;
                seen[u] = -1;
                cycs.push_front({u, time, {&Q[qi], &Q[end]}});
            }
        }
        for (int i = 0; i < qi; i++) in[dsu.getRoot(Q[i].v)] = Q[i];
    }

    for (auto& [u, t, comp] : cycs) {
        dsu.rollback(t);
        Edge inEdge = in[u];
        for (auto& e : comp) in[dsu.getRoot(e.v)] = e;
        in[dsu.getRoot(inEdge.v)] = inEdge;
    }

    vector<int> par(n);
    for (int i = 0; i < n; i++) par[i] = in[i].u;
    return {res, par};
}
#line 7 "Graph/tests/directed_mst.test.cpp"

#define REP(i, a) for (int i = 0, _##i = (a); i < _##i; ++i)

int32_t main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n, m, root; cin >> n >> m >> root;
    vector<Edge> edges(m);
    for (auto& e : edges) {
        cin >> e.u >> e.v >> e.cost;
    }
    auto [total, par] = directed_mst(n, root, edges);
    cout << total << endl;
    REP(i,n) cout << (i == root ? root : par[i]) << ' ';
    cout << endl;
    return 0;
}