Mst oriented in O(NM)

algo/mst_oriented/01.in

@@ -0,0 +1,6 @@
+5 5 0
+0 1 5
+0 2 3
+2 3 2
+3 4 2
+4 2 2

algo/mst_oriented/02.in

@@ -0,0 +1,12 @@
+9 11 0
+0 1 4
+2 6 4
+1 2 3
+2 3 3
+3 4 3
+4 5 3
+5 1 3
+4 7 1
+7 5 1
+5 4 1
+7 8 2

algo/mst_oriented/mst_oriented.cpp

@@ -0,0 +1,202 @@
+#include <bits/stdc++.h>
+
+#define FOR(i, n) for(lli i = 0; i < (lli)(n); ++i)
+
+#define X(A) get<0>(A)
+#define Y(A) get<1>(A)
+#define Z(A) get<2>(A)
+#define W(A) get<3>(A)
+
+#define mt make_tuple
+
+using namespace std;
+using lli = long long int;
+
+using pii   = tuple<lli, lli>;
+using piii  = tuple<lli, lli, lli>;
+using vi    = vector<lli>;
+using vii   = vector<pii>;
+using viii  = vector<piii>;
+using vvi   = vector<vi>;
+using vvii  = vector<vii>;
+using vviii = vector<viii>;
+using vb    = vector<bool>;
+using vvb   = vector<vb>;
+
+/** WARNING **
+ *************
+
+ This algorithm considers the graph to have no self-pointing edge, ie. no edge
+ (u, u)
+*/
+
+const int INFTY = 1000*1000*1000; // FIXME large enough for this problem?
+
+struct Edge {
+    int from, to;
+    int weight;
+    Edge() {}
+    Edge(int from, int to, int weight): from(from), to(to), weight(weight) {}
+};
+
+struct Vert {
+    int selfId;
+    int ufFather;
+    int ufHeight;
+    vector<int> out;
+    vector<int> in;
+    int parentEdge;
+    bool suppressed;
+
+    Vert(int id):
+        selfId(id), ufFather(id), ufHeight(1),
+        parentEdge(-1), suppressed(false)
+    {}
+};
+
+int root;
+vector<Edge> edges;
+vector<Vert> vertices;
+int cost = 0;
+
+int ufFind(int vert) {
+    Vert& v = vertices[vert];
+    if(v.ufFather == vert)
+        return vert;
+    v.ufFather = ufFind(v.ufFather);
+    return v.ufFather;
+}
+
+void ufUnion(int iv1, int iv2) {
+    iv1 = ufFind(iv1);
+    iv2 = ufFind(iv2);
+    Vert& v1 = vertices[iv1];
+    Vert& v2 = vertices[iv2];
+    if(v1.ufHeight < v2.ufHeight)
+        v1.ufFather = iv2;
+    else if(v2.ufHeight < v1.ufHeight)
+        v2.ufFather = iv1;
+    else {
+        v1.ufFather = iv2;
+        v2.ufHeight++;
+    }
+}
+
+void buildMst(int vertId) {
+    printf("> %d\n", vertId);
+    Vert& vert = vertices[vertId];
+    int minIncWeight = INFTY;
+    int minId = -1;
+
+    for(int incEdge: vert.in) {
+        Edge& edge = edges[incEdge];
+        if(minIncWeight > edge.weight) {
+            minIncWeight = edge.weight;
+            minId = incEdge;
+        }
+    }
+
+    if(minId == -1) {
+        assert(vertId == root);
+        return;
+    }
+
+    for(int incEdge: vert.in) {
+        Edge& edge = edges[incEdge];
+        edge.weight -= minIncWeight;
+    }
+    cost += minIncWeight;
+
+    vert.parentEdge = minId;
+    int parent = edges[minId].from;
+    int parentCmp = ufFind(parent), selfCmp = ufFind(vertId);
+    if(parentCmp == selfCmp) {
+        // Build current cycle
+        vector<int> cycle;
+        cycle.push_back(parent);
+        while(cycle.back() != vertId)
+            cycle.push_back(edges[vertices[cycle.back()].parentEdge].from);
+
+        vertices.push_back(Vert(vertices.size()));
+        Vert& condensed = vertices.back();
+        printf("Compressing into %d: ", condensed.selfId);
+        for(int cyc: cycle)
+            printf("%d ", cyc);
+        puts("");
+        vector<int> nInEdges, nOutEdges;
+        for(int cycleVert: cycle) {
+            Vert& v = vertices[cycleVert];
+            v.suppressed = true;
+            for(int inEdge: v.in) {
+                edges[inEdge].to = condensed.selfId;
+                nInEdges.push_back(inEdge);
+            }
+            for(int outEdge: v.out) {
+                edges[outEdge].from = condensed.selfId;
+                if(edges[outEdge].to != condensed.selfId)
+                    condensed.out.push_back(outEdge);
+            }
+        }
+
+        for(int inEdge: nInEdges) {
+            if(edges[inEdge].from != condensed.selfId)
+                condensed.in.push_back(inEdge);
+        }
+
+        ufUnion(vertId, condensed.selfId);
+        buildMst(condensed.selfId);
+    }
+    else {
+        ufUnion(vertId, parent);
+        if(ufFind(vertId) != ufFind(root))
+            buildMst(parent);
+    }
+}
+
+void dfs(int pos, vector<bool>& seen) {
+    if(seen[pos])
+        return;
+    seen[pos] = true;
+    for(int outEdge: vertices[pos].out)
+        dfs(edges[outEdge].to, seen);
+}
+
+bool allReachable() {
+    vector<bool> seen(vertices.size(), false);
+    dfs(root, seen);
+    for(size_t pos=0; pos < seen.size(); ++pos) {
+        if(!seen[pos]) {
+            fprintf(stderr, "NOT REACHABLE: %lu\n", pos);
+            return false;
+        }
+    }
+    return true;
+}
+
+int main() {
+    int nbVert, nbEdge;
+    scanf("%d%d%d", &nbVert, &nbEdge, &root);
+    for(int vert=0; vert < nbVert; ++vert)
+        vertices.push_back(Vert(vert));
+    for(int edgeId=0; edgeId < nbEdge; ++edgeId) {
+        Edge edge;
+        scanf("%d%d%d", &edge.from, &edge.to, &edge.weight);
+        edges.push_back(edge);
+        vertices[edge.from].out.push_back(edges.size() - 1);
+        vertices[edge.to].in.push_back(edges.size() - 1);
+    }
+
+    if(!allReachable()) {
+        printf("-1\n");
+        puts("BLEH");
+        return 0;
+    }
+
+    for(int vert=0; vert < nbVert; ++vert) {
+        if(vertices[vert].suppressed || ufFind(vert) == root)
+            continue;
+        buildMst(vert);
+    }
+
+    printf("%d\n", cost);
+}