Added directedMST/Minimum arborescence code (#110)

Author: Chillee

content/data-structures/UnionFind.h

@@ -14,10 +14,11 @@ struct UF {
 	bool same_set(int a, int b) { return find(a) == find(b); }
 	int size(int x) { return -e[find(x)]; }
 	int find(int x) { return e[x] < 0 ? x : e[x] = find(e[x]); }
-	void join(int a, int b) {
+	bool join(int a, int b) {
 		a = find(a), b = find(b);
-		if (a == b) return;
+		if (a == b) return false;
 		if (e[a] > e[b]) swap(a, b);
 		e[a] += e[b]; e[b] = a;
+		return true;
 	}
 };

content/graph/DirectedMST.h

@@ -0,0 +1,62 @@
+/**
+ * Author: chilli, Takanori MAEHARA
+ * 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
+ * tree/arborescence of a directed graph, given a root node. If no MST exists, returns -1.
+ * Time: O(E \log V)
+ * Status: Fuzz-tested, also tested on NWERC 2018 fastestspeedrun
+ */
+#pragma once
+
+#include "../data-structures/UnionFind.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); }
+
+ll dmst(int n, int r, vector<Edge>& g) {
+	UF uf(n);
+	vector<Node*> heap(n);
+	trav(e, g) heap[e.b] = merge(heap[e.b], new Node{e});
+	ll res = 0;
+	vi seen(n, -1), path(n);
+	seen[r] = r;
+	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;
+			Edge e = heap[u]->top();
+			heap[u]->delta -= e.w, pop(heap[u]);
+			res += e.w, u = uf.find(e.a);
+			if (seen[u] == s) {
+				Node* cyc = 0;
+				do cyc = merge(cyc, heap[w = path[--qi]]);
+				while (uf.join(u, w));
+				u = uf.find(u);
+				heap[u] = cyc, seen[u] = -1;
+			}
+		}
+	}
+	return res;
+}

content/graph/chapter.tex

@@ -39,5 +39,6 @@ \section{Trees}
 	\kactlimport{CompressTree.h}
 	\kactlimport{HLD.h}
 	\kactlimport{LinkCutTree.h}
+	\kactlimport{DirectedMST.h}
 	\kactlimport{MatrixTree.h}
-	\columnbreak
+	\hardcolumnbreak

fuzz-tests/graph/DirectedMST.cpp

@@ -0,0 +1,169 @@
+#include <bits/stdc++.h>
+using namespace std;
+
+#define rep(i, a, b) for(int i = a; i < int(b); ++i)
+#define trav(a, v) for(auto& a : v)
+#define all(x) x.begin(), x.end()
+#define sz(x) (int)(x).size()
+
+typedef long long ll;
+typedef pair<int, int> pii;
+typedef vector<int> vi;
+
+struct Bumpalloc {
+	char buf[450 << 20];
+	size_t bufp;
+	void* alloc(size_t s) {
+		assert(s < bufp);
+		return (void*)&buf[bufp -= s];
+	}
+	Bumpalloc() { reset(); }
+
+	template<class T> T* operator=(T&& x) {
+		T* r = (T*)alloc(sizeof(T));
+		new(r) T(move(x));
+		return r;
+	}
+	void reset() { bufp = sizeof buf; }
+} bumpalloc;
+
+// When not testing perf, we don't want to leak memory
+#ifndef TEST_PERF
+#define new bumpalloc =
+#endif
+#include "../../content/graph/DirectedMST.h"
+#ifndef TEST_PERF
+#undef new
+#endif
+
+namespace mit {
+
+#define N 110000
+#define M 110000
+#define inf 2000000000
+
+struct edg {
+    int u, v;
+    int cost;
+} E[M], E_copy[M];
+
+int In[N], ID[N], vis[N], pre[N];
+
+// edges pointed from root.
+int Directed_MST(int root, int NV, int NE) {
+    for (int i = 0; i < NE; i++)
+        E_copy[i] = E[i];
+    int ret = 0;
+    int u, v;
+    while (true) {
+        rep(i,0,NV)   In[i] = inf;
+        rep(i,0,NE) {
+            u = E_copy[i].u;
+            v = E_copy[i].v;
+            if(E_copy[i].cost < In[v] && u != v) {
+                In[v] = E_copy[i].cost;
+                pre[v] = u;
+            }
+        }
+        rep(i,0,NV) {
+            if(i == root)   continue;
+            if(In[i] == inf)    return -1; // no solution
+        }
+
+        int cnt = 0;
+        rep(i,0,NV) {
+            ID[i] = -1;
+            vis[i] = -1;
+        }
+        In[root] = 0;
+
+        rep(i,0,NV) {
+            ret += In[i];
+            int v = i;
+            while(vis[v] != i && ID[v] == -1 && v != root) {
+                vis[v] = i;
+                v = pre[v];
+            }
+            if(v != root && ID[v] == -1) {
+                for(u = pre[v]; u != v; u = pre[u]) {
+                    ID[u] = cnt;
+                }
+                ID[v] = cnt++;
+            }
+        }
+        if(cnt == 0)    break;
+        rep(i,0,NV) {
+            if(ID[i] == -1) ID[i] = cnt++;
+        }
+        rep(i,0,NE) {
+            v = E_copy[i].v;
+            E_copy[i].u = ID[E_copy[i].u];
+            E_copy[i].v = ID[E_copy[i].v];
+            if(E_copy[i].u != E_copy[i].v) {
+                E_copy[i].cost -= In[v];
+            }
+        }
+        NV = cnt;
+        root = ID[root];
+    }
+    return ret;
+}
+}
+
+int adj[105][105];
+int main() {
+	rep(it,0,100000) {
+		bumpalloc.reset();
+		int n = (rand()%20)+1;
+		int density = rand() % 101;
+		int r = rand()%n;
+		int cnt = 0;
+		vector<Edge> edges;
+		rep(i,0,n)
+			rep(j,0,n){
+				if (i==j) continue;
+				if (rand() % 100 >= density) continue;
+				int weight = rand()%100;
+				mit::E[cnt++] = {i,j, weight};
+				edges.push_back({i,j,weight});
+				adj[i][j] = weight;
+			}
+		ll ans1 = mit::Directed_MST(r, n, cnt);
+		ll ans2 = dmst(n, r, edges);
+		assert(ans1 == ans2);
+		// For verifying a reconstruction:
+		/*
+		if (ans1 != -1) {
+			vi par = pa.second;
+			if (debug) {
+				cout << "r = " << r << endl;
+				trav(x, par) cout << x << ' ';
+				cout << endl;
+				trav(e, edges) {
+					cout << e.a << ' ' << e.b << ' ' << e.w << endl;
+				}
+			}
+			ll sum = 0;
+			vector<vi> ch(n);
+			rep(i,0,n) {
+				if (i == r) assert(par[i] == -1);
+				else {
+					assert(par[i] != -1);
+					sum += adj[par[i]][i];
+					ch[par[i]].push_back(i);
+				}
+			}
+			assert(sum == ans1);
+			vi seen(n), q = {r};
+			rep(qi,0,sz(q)) {
+				int s = q[qi];
+				if (!seen[s]++)
+					trav(x, ch[s]) q.push_back(x);
+			}
+			assert(count(all(seen), 0) == 0);
+		}
+		*/
+	}
+	cout<<"Tests passed!"<<endl;
+	return 0;
+}