CP Notebook

← all categories

Mcmf b9362c4f

Min-cost max-flow via Dijkstra with Johnson's potentials (a PBDS pairing heap gives O(1) decrease-key). If costs can be negative, call setpi(s) once before maxflow (negative cost cycles are not supported). Add every edge via addEdge() before the first maxflow() call — edges are stored by value per-node and augmenting paths hold raw pointers into them, which addEdge()'s push_back could invalidate mid-flow.

Time: O(FE \log V), F = total flow; O(VE) for setpi. kactl-derived, tested

content/flows/mcmf.h

#undef int  // pb_ds headers use 'int' internally; #define int long long corrupts them
#include <ext/pb_ds/priority_queue.hpp>
#define int long long

template <typename T>
struct MCMF {
    struct Edge { int from, to, rev; T cap, cost, flow; };
    int n;
    vector<vector<Edge>> adj;
    vector<bool> seen;
    vector<T> dist, pi;
    vector<Edge *> par;
    T INF = numeric_limits<T>::max() / 4;

    MCMF(int n) : n(n), adj(n), seen(n), dist(n), pi(n, T()), par(n) {}

    void addEdge(int from, int to, T cap, T cost) {
        if (from == to) return;
        adj[from].push_back({from, to, (int)adj[to].size(), cap, cost, T()});
        adj[to].push_back({to, from, (int)adj[from].size() - 1, T(), -cost, T()});
    }

    void path(int s) {
        fill(seen.begin(), seen.end(), false);
        fill(dist.begin(), dist.end(), INF);
        dist[s] = T();
        __gnu_pbds::priority_queue<pair<T, int>> q;
        vector<typename decltype(q)::point_iterator> its(n, q.end());
        q.push({T(), s});
        while (!q.empty()) {
            s = q.top().second; q.pop();
            seen[s] = true;
            T di = dist[s] + pi[s];
            for (Edge &e : adj[s]) {
                if (seen[e.to]) continue;
                T val = di - pi[e.to] + e.cost;
                if (e.cap - e.flow > 0 && val < dist[e.to]) {
                    dist[e.to] = val;
                    par[e.to] = &e;
                    if (its[e.to] == q.end()) its[e.to] = q.push({-dist[e.to], e.to});
                    else q.modify(its[e.to], {-dist[e.to], e.to});
                }
            }
        }
        for (int i = 0; i < n; i++) pi[i] = min(pi[i] + dist[i], INF);
    }

    pair<T, T> maxflow(int s, int t) {
        T totalFlow = T(), totalCost = T();
        for (path(s); seen[t]; path(s)) {
            T f = INF;
            for (Edge *e = par[t]; e; e = par[e->from]) f = min(f, e->cap - e->flow);
            totalFlow += f;
            for (Edge *e = par[t]; e; e = par[e->from]) {
                e->flow += f;
                adj[e->to][e->rev].flow -= f;
            }
        }
        for (int i = 0; i < n; i++) for (Edge &e : adj[i]) totalCost += e.cost * e.flow;
        return {totalFlow, totalCost / 2};
    }

    // call before maxflow() only if some costs are negative
    void setpi(int s) {
        fill(pi.begin(), pi.end(), INF);
        pi[s] = T();
        int rounds = n, changed = 1;
        while (changed-- && rounds--) {
            for (int i = 0; i < n; i++) {
                if (pi[i] == INF) continue;
                for (Edge &e : adj[i]) {
                    if (e.cap == 0) continue;
                    if (pi[i] + e.cost < pi[e.to]) { pi[e.to] = pi[i] + e.cost; changed = 1; }
                }
            }
        }
        assert(rounds >= 0);  // negative cost cycle
    }
};