-
Notifications
You must be signed in to change notification settings - Fork 700
/
DirectedMST.h
75 lines (71 loc) · 2.19 KB
/
DirectedMST.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
/**
* Author: chilli, Takanori MAEHARA, Benq, Simon Lindholm
* Date: 2019-05-10
* License: CC0
* Source: https://github.com/spaghetti-source/algorithm/blob/master/graph/arborescence.cc
* and https://github.com/bqi343/USACO/blob/42d177dfb9d6ce350389583cfa71484eb8ae614c/Implementations/content/graphs%20(12)/Advanced/DirectedMST.h for the reconstruction
* Description: Finds a minimum spanning
* tree/arborescence of a directed graph, given a root node. If no MST exists, returns -1.
* Time: O(E \log V)
* Status: Stress-tested, also tested on NWERC 2018 fastestspeedrun
*/
#pragma once
#include "../data-structures/UnionFindRollback.h"
struct Edge { int a, b; ll w; };
struct Node { /// lazy skew heap node
Edge key;
Node *l, *r;
ll delta;
void prop() {
key.w += delta;
if (l) l->delta += delta;
if (r) r->delta += delta;
delta = 0;
}
Edge top() { prop(); return key; }
};
Node *merge(Node *a, Node *b) {
if (!a || !b) return a ?: b;
a->prop(), b->prop();
if (a->key.w > b->key.w) swap(a, b);
swap(a->l, (a->r = merge(b, a->r)));
return a;
}
void pop(Node*& a) { a->prop(); a = merge(a->l, a->r); }
pair<ll, vi> dmst(int n, int r, vector<Edge>& g) {
RollbackUF uf(n);
vector<Node*> heap(n);
for (Edge e : g) heap[e.b] = merge(heap[e.b], new Node{e});
ll res = 0;
vi seen(n, -1), path(n), par(n);
seen[r] = r;
vector<Edge> Q(n), in(n, {-1,-1}), comp;
deque<tuple<int, int, vector<Edge>>> cycs;
rep(s,0,n) {
int u = s, qi = 0, w;
while (seen[u] < 0) {
if (!heap[u]) return {-1,{}};
Edge e = heap[u]->top();
heap[u]->delta -= e.w, pop(heap[u]);
Q[qi] = e, path[qi++] = u, seen[u] = s;
res += e.w, u = uf.find(e.a);
if (seen[u] == s) { /// found cycle, contract
Node* cyc = 0;
int end = qi, time = uf.time();
do cyc = merge(cyc, heap[w = path[--qi]]);
while (uf.join(u, w));
u = uf.find(u), heap[u] = cyc, seen[u] = -1;
cycs.push_front({u, time, {&Q[qi], &Q[end]}});
}
}
rep(i,0,qi) in[uf.find(Q[i].b)] = Q[i];
}
for (auto& [u,t,comp] : cycs) { // restore sol (optional)
uf.rollback(t);
Edge inEdge = in[u];
for (auto& e : comp) in[uf.find(e.b)] = e;
in[uf.find(inEdge.b)] = inEdge;
}
rep(i,0,n) par[i] = in[i].a;
return {res, par};
}