Directed MST reconstruction (kth-competitive-programming#153)

* UnionFind with rollback

* Directed MST reconstruction

* Include UnionFindRollback in the pdf, remove UnionFind

Its functionality is a superset, and one should know UnionFind by heart
anyway.

content/data-structures/UnionFindRollback.h

@@ -0,0 +1,33 @@
+/**
+ * Author: Lukas Polacek, Simon Lindholm
+ * Date: 2019-12-26
+ * License: CC0
+ * Source: folklore
+ * Description: Disjoint-set data structure with undo.
+ * If undo is not needed, skip st, time() and rollback().
+ * Usage: int t = uf.time(); ...; uf.rollback(t);
+ * Time: $O(\log(N))$
+ */
+#pragma once
+
+struct RollbackUF {
+	vi e; vector<pii> st;
+	RollbackUF(int n) : e(n, -1) {}
+	int size(int x) { return -e[find(x)]; }
+	int find(int x) { return e[x] < 0 ? x : find(e[x]); }
+	int time() { return sz(st); }
+	void rollback(int t) {
+		for (int i = time(); i --> t;)
+			e[st[i].first] = st[i].second;
+		st.resize(t);
+	}
+	bool join(int a, int b) {
+		a = find(a), b = find(b);
+		if (a == b) return false;
+		if (e[a] > e[b]) swap(a, b);
+		st.push_back({a, e[a]});
+		st.push_back({b, e[b]});
+		e[a] += e[b]; e[b] = a;
+		return true;
+	}
+};

content/data-structures/chapter.tex

@@ -4,7 +4,8 @@ \chapter{Data structures}
 \kactlimport{HashMap.h}
 \kactlimport{SegmentTree.h}
 \kactlimport{LazySegmentTree.h}
-\kactlimport{UnionFind.h}
+% \kactlimport{UnionFind.h}
+\kactlimport{UnionFindRollback.h}
 \kactlimport{SubMatrix.h}
 \kactlimport{Matrix.h}
 \kactlimport{LineContainer.h}

content/graph/DirectedMST.h

@@ -1,16 +1,17 @@
 /**
- * Author: chilli, Takanori MAEHARA
+ * 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
- * Description: Edmonds' algorithm for finding the weight of the minimum spanning
+ * 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/UnionFind.h"
+#include "../data-structures/UnionFindRollback.h"
 
 struct Edge { int a, b; ll w; };
 struct Node { /// lazy skew heap node
@@ -34,29 +35,41 @@ Node *merge(Node *a, Node *b) {
 }
 void pop(Node*& a) { a->prop(); a = merge(a->l, a->r); }
 
-ll dmst(int n, int r, vector<Edge>& g) {
-	UF uf(n);
+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);
+	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) {
-			path[qi++] = u, seen[u] = s;
-			if (!heap[u]) return -1;
+			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) {
+			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;
+				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];
 	}
-	return res;
+
+	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};
 }

stress-tests/graph/DirectedMST.cpp

@@ -102,7 +102,7 @@ int Directed_MST(int root, int NV, int NE) {
 
 int adj[105][105];
 int main() {
-	rep(it,0,100000) {
+	rep(it,0,50000) {
 		bumpalloc.reset();
 		int n = (rand()%20)+1;
 		int density = rand() % 101;
@@ -118,14 +118,16 @@ int main() {
 				edges.push_back({i,j,weight});
 				adj[i][j] = weight;
 			}
+
 		ll ans1 = mit::Directed_MST(r, n, cnt);
-		ll ans2 = dmst(n, r, edges);
+		auto pa = dmst(n, r, edges);
+		ll ans2 = pa.first;
 		assert(ans1 == ans2);
-		// For verifying a reconstruction:
-		/*
+
+		// Verifying reconstruction:
 		if (ans1 != -1) {
 			vi par = pa.second;
-			if (debug) {
+			if (0) {
 				cout << "r = " << r << endl;
 				for(auto &x: par) cout << x << ' ';
 				cout << endl;
@@ -152,7 +154,6 @@ int main() {
 			}
 			assert(count(all(seen), 0) == 0);
 		}
-		*/
 	}
 	cout<<"Tests passed!"<<endl;
 	return 0;