The 2nd Universal Cup. Stage 2: SPb K. Poor Students

問題のリンク：https://qoj.ac/contest/1356/problem/7185

問題概要

$N$ 人の生徒と $K$ 個のコースがあり、各生徒は1つのコースを受講する必要がある。

$i$ 番目の生徒が、$j$ 番目のコースを受講する不満度は $c _ {i,j}$ である。

ただし、$i$ 番目のコースは高々 $a _ {i}$ 人のみ受講可能である。

この時、生徒全体のストレスの総和を最小化せよ。

制約

$1 \le N \le 50000$
$1 \le K \le 10$

解法

一旦制約を無視して考えると、生徒側とコース側の頂点に分けた単純な最小費用流問題として記述できることが分かります。~~ちなみにバカ早い cost-scaling Push-Relabel を窃盗してくるとこの時点でAC出来てしまう…*1。~~

ここで最小費用流のアルゴリズムである、最短路反復の気持ちになって考察を進めます。アルゴリズムの解説はhosさんのスライドが詳しいです。

要するに最短路を1つ見つけ、その最短路に従って残余グラフを費用付きに更新ということを流用分だけ繰り返していけば、最小費用流が求まるという内容です。

ただしナイーブに本問題に適用しようとすると、1回のフロー毎に頂点数 $N+K$、辺数 $NK$ のグラフを扱う必要があり間に合いません。そこで、シンク側の頂点数が少ないことを利用したアルゴリズムの改善を行います。

最短路反復で得られる最短路は、生徒 $i _ 1$ → コース $j _ 1$ → 生徒 $i _ 2$ → コース $j _ 2$ → … → コース $j _ k$ となっていること、コース間の移動は最も短い辺のみ考慮すれば良いことを踏まえると、頂点数 $K+1$ のグラフに圧縮できそうという着想が得られます。$+1$ としているのは、生徒側から伸びる未使用の辺を1つの頂点に集約して表現するためです。

1度のフロー毎に上記グラフの作成を行えば、通常の最短路反復のアルゴリズムに合流します。グラフ作成のための差分更新は、生徒側の頂点が今どのコースとマッチングしているかの情報を持ち、辺集合を priority_queue などで管理することで可能です。

以上を実装することでAC可能です。詳しい実装方針や計算量はソースコードを参照してください。

ちなみに HoMaMa はこの問題を開始6分でACしているが、チームメイトのうち2人が上記問題を想定計算量で解くライブラリを持っていたらしい*2。

実装

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;

mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
ll myRand(ll B) { return (ull)rng() % B; }

// m sources, supply 1
// n sinks, capacity cap[v]
// calc min cost maxflow : O(m n^3 + m n^2 log m)
template <typename cost_t> struct MCF {
    int m, n;
    vector<int> cap, to; // to : source 側の頂点がどの sink と対応しているか
    vector<vector<cost_t>> cost;

    using P = pair<cost_t, int>;
    const cost_t inf = numeric_limits<cost_t>::max() / 2;
    vector<priority_queue<P, vector<P>, greater<P>>> edge;
    vector<cost_t> adj, dist;
    vector<int> prv;
    vector<bool> on;

    MCF(int m, int n) : m(m), n(n), cap(n), cost(m, vector<cost_t>(n + 1)) {
        to = vector<int>(m, n);
        edge.resize((n + 1) * n);
        adj.resize((n + 1) * n);
        dist.resize(n);
        prv.resize(n);
        on.resize(n);
    }
    void set_cap(int i, int c) { cap[i] = c; }
    void set_cost(int i, int j, cost_t x) {
        cost[i][j] = x;
        edge[n * n + j].push({cost[i][j], i});
    }

    cost_t solve() {
        cost_t res = 0;
        for (int rep = 0; rep < m; ++rep) {
            // sink 側のグラフ構築
            for (int i = 0; i <= n; ++i) {
                for (int j = 0; j < n; ++j) {
                    auto &pq = edge[i * n + j];
                    while (pq.size()) {
                        if (to[pq.top().second] == i) break;
                        else pq.pop();
                    }
                    adj[i * n + j] = (pq.empty() ? inf : pq.top().first);
                }
            }
            // 最短路の計算
            queue<int> q;
            for (int i = 0; i < n; ++i) {
                dist[i] = adj[n * n + i], prv[i] = n, on[i] = true;
                q.push(i);
            }
            while (q.size()) {
                int s = q.front();
                q.pop();
                on[s] = false;
                for (int t = 0; t < n; ++t) {
                    cost_t c = dist[s] + adj[s * n + t];
                    if (dist[t] > c) {
                        dist[t] = c, prv[t] = s;
                        if (!on[t]) {
                            on[t] = true;
                            q.push(t);
                        }
                    }
                }
            }
            // 求めた最短路の1つを復元してフローを更新する
            int s = -1;
            for (int i = 0; i < n; ++i) {
                if (cap[i] > 0) {
                    if (s == -1) s = i;
                    if (dist[i] < dist[s]) s = i;
                }
            }
            assert(s != -1);
            res += dist[s], cap[s] -= 1;
            while (s != n) {
                int t = prv[s], idx = edge[t * n + s].top().second;
                for (int r = 0; r < n; ++r) {
                    if (s != r) {
                        edge[s * n + r].push({cost[idx][r] - cost[idx][s], idx});
                    }
                }
                to[idx] = s, s = t;
            }
        }
        return res;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, k; cin >> n >> k;
    MCF<ll> mcf(n, k);
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < k; ++j) {
            ll cost; cin >> cost;
            mcf.set_cost(i, j, cost);
        }
    }
    for (int i = 0; i < k; ++i) {
        int cap; cin >> cap;
        mcf.set_cap(i, cap);
    }
    cout << mcf.solve() << endl;
}

*1:ecnerwala が言及していた窃盗先 https://github.com/dacin21/dacin21_codebook/blob/master/flow/mincost_PRonly.cpp

*2:nowcoder's multi-uni で既出とのこと

かっつのメモ帳

主に競プロ　時々日記

The 2nd Universal Cup. Stage 2: SPb K. Poor Students

問題概要

解法

実装