MinCostFlow

gistfile1.cpp

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>
#include <set>
 
#define loop(n, i) for(int i=0;i<n;i++)
#define from_to(from, to, i) for(int i=from;i<=to;i++)

using namespace std;

using P = pair<int, int>;
class MCF
{
public:
    static const int MAX_V = 10000;
    const int INF = 1<<29;
    int V;
    struct edge { int to, cap, cost, rev; };

    vector<edge> G[MAX_V];
    int H[MAX_V];
    int D[MAX_V];
    int prevV[MAX_V];
    int prevE[MAX_V];

    void add_edge(int from, int to, int cap, int cost) {
        G[from].push_back(edge { to, cap, cost, (int)G[to].size() });
        G[to].push_back(edge { from, 0, -cost, (int)G[from].size()-1 });
    }

    int calc(int s, int t, int f) {
        int res = 0;
        fill(H, H+V, 0);
        while (f > 0) {
            priority_queue<P, vector<P>, greater<P>> q;
            fill(D, D+V, INF);
            D[s] = 0;
            q.emplace(0, s);
            while (!q.empty()) {
                P p = q.top(); q.pop();
                int v = p.second;
                if (D[v] < p.first) continue;
                loop (G[v].size(), i) {
                    edge e = G[v][i];
                    if (e.cap > 0 && D[e.to] > D[v] + e.cost + H[v] - H[e.to]) {
                        D[e.to] = D[v] + e.cost + H[v] - H[e.to];
                        prevV[e.to] = v;
                        prevE[e.to] = i;
                        q.emplace(D[e.to], e.to);
                    }
                }
            }
            if (D[t] == INF) return -1;
            loop (V, v) H[v] += D[v];
            int d = f;
            for (int v = t; v != s; v = prevV[v]) {
                d = min(d, G[prevV[v]][prevE[v]].cap);
            }
            f -= d;
            res += d * H[t];
            for (int v = t; v != s; v = prevV[v]) {
                edge &e = G[prevV[v]][prevE[v]];
                e.cap -= d;
                G[v][e.rev].cap += d;
            }
        }
        return res;
    }
};

class SpecialCells
{
public:
    int guess (vector <int> x, vector <int> y)
    {
        vector<int> xs, ys;
        for (int c : set<int>(x.begin(), x.end())) xs.push_back(c);
        for (int &c : x) c = distance(xs.begin(), find(xs.begin(), xs.end(), c));
        for (int c : set<int>(y.begin(), y.end())) ys.push_back(c);
        for (int &c : y) c = distance(ys.begin(), find(ys.begin(), ys.end(), c));
        int xn = xs.size(), yn = ys.size();

        vector<int> cntX(xn), cntY(yn);
        for (int c : x) cntX[c]++;
        for (int c : y) cntY[c]++;

        MCF mcf;
        mcf.V = xn + yn + 2;
        int src = xn+yn, sink = src+1;

        auto spc = [&](int cx, int cy) -> int {
            loop (x.size(), i) if (cx == x[i] && cy == y[i]) return 1;
            return 0;
        };

        loop (xn, i) mcf.add_edge(src, i, cntX[i], 0);
        loop (xn, xi) loop (yn, yi) mcf.add_edge(xi, xn+yi, 1, spc(xi, yi));
        loop (yn, i) mcf.add_edge(xn+i, sink, cntY[i], 0);
        return mcf.calc(src, sink, x.size());
    }
};