けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 326 G - Unlock Achievement (橙色, 625 点)

2 変数劣モジュラ関数の和の最小化を最小カットにするやつ!

問題概要

 N 種類のスキルがある。それぞれ初期状態のスキルレベルは 1 であるが、最大で 5 まで上げられる。スキル  i のスキルレベルを 1 上げるのに必要なコストは  C_{i} 円である。

 M 種類のアチーブメントがある。アチーブメント  i は、各スキル  j のスキルレベルを  L_{i, j} 以上にすることで達成される。達成されると  A_{i} 円の報酬がもらえる。

制約

  •  1 \le N, M \le 50
  •  1 \le L_{i, j} \le 5

考えたこと

いかにも劣モジュラ関数っぽい問題設定である。スキルレベルが 5 段階あるのが厄介にも思えるが、5 値を取ると思われる変数も、4 個の 0-1 変数で表すことができる。次のようにする。


  •  x_{i, 0} ← スキル  i のレベルが 2 以上であるとき True (そうでないとき False)
  •  x_{i, 1} ← スキル  i のレベルが 3 以上であるとき True (そうでないとき False)
  •  x_{i, 2} ← スキル  i のレベルが 4 以上であるとき True (そうでないとき False)
  •  x_{i, 3} ← スキル  i のレベルが 5 以上であるとき True (そうでないとき False)

こうすると、たとえばスキルレベルが 1 であることは (False, False, False, False) と表せるし、スキルレベルが 3 であることは (True, True, False, False) と表せるし、スキルレベルが 5 であることは (True, True, True, True) と表せる。

そして、各変数  x_{i, j} に対して、次のようなコストを考慮すればよい。


  •  x_{i, j} = True ならば  C_{i} のコストがかかり、 x_{i, j} = False ならば  0 のコストがかかる (1 変数劣モジュラ)
  •  j \lt k に対して、 x_{i, j} = False かつ  x_{i, k} = True であるとき、 \infty のペナルティを与える (2 変数劣モジュラ)

アチーブメントについて

アチーブメント制約は「 N 個の変数がすべて True ならば、ある利得が得られる」という形をしている。このような「すべて True ならば利得」や「すべて False ならば利得」といったものも、最小カットで表現できることがよく知られている。

ここでは、 N 個の変数  z_{0}, z_{1}, \dots, z_{N-1} がすべて True のときに、利得  G が得られるとしよう。このとき、具体的には新たに次の変数  y を用意すればよい。


  • 「変数  y:利得  G が得られるならば True (得られないならば False)」とする
  •  y = True ならば  -G のコストがかかり、 y = False ならば  0 のコストがかかる (1 変数劣モジュラ)
  • 各変数  z_{i} について、 z_{i} = False かつ  y = True であるとき、 \infty のペナルティを課す (2 変数劣モジュラ)

コード

以上の 2 変数劣モジュラ関数の和のグラフ表現をほぼ機械的に処理できるようなライブラリを整備した。その際には、「 N 個の変数がすべて True ならば、ある利得が得られる」という部分についてもライブラリに含めた。

グラフのサイズは頂点数が  4N + M 程度であり、計算時間には余裕がある。

#include <bits/stdc++.h>
using namespace std;


/*
 N 個の bool 変数 x_0, x_1, ..., x_{N-1} について、以下の形のコストが定められたときの最小コストを求める
 
 ・1 変数 xi に関するコスト (1 変数劣モジュラ関数)
    xi = F のときのコスト, xi = T のときのコスト
 
 ・2 変数 xi, xj 間の関係性についてのコスト (2 変数劣モジュラ関数)
   (xi, xj) = (F, F): コスト A
   (xi, xj) = (F, T): コスト B
   (xi, xj) = (T, F): コスト C
   (xi, xj) = (T, T): コスト D
  (ただし、B + C >= A + D でなければならない)
 
 ・よくある例は、A = B = D = 0, C >= 0 の形である (特に関数化している)
    ・この場合は、特に Project Selection Problem と呼ばれ、ローカルには「燃やす埋める」などとも呼ばれる
 ・他に面白い例として、A = B = C = 0, D < 0 の形もある (AOJ 2903 Board)

 */


// 2-variable submodular optimization
template<class COST> struct TwoVariableSubmodularOpt {
    // edge class
    struct Edge {
        // core members
        int rev, from, to;
        COST cap, icap, flow;
        
        // constructor
        Edge(int r, int f, int t, COST c)
        : rev(r), from(f), to(t), cap(c), icap(c), flow(0) {}
        void reset() { cap = icap, flow = 0; }
        
        // debug
        friend ostream& operator << (ostream& s, const Edge& E) {
            return s << E.from << "->" << E.to << '(' << E.flow << '/' << E.icap << ')';
        }
    };
    
    // inner data
    int N, S, T;
    COST OFFSET, INF;
    vector<vector<Edge>> list;
    vector<pair<int,int>> pos;
    
    // constructor
    TwoVariableSubmodularOpt() : N(2), S(0), T(0), OFFSET(0) {}
    TwoVariableSubmodularOpt(int n, COST inf = 0)
    : N(n), S(n), T(n + 1), OFFSET(0), INF(inf), list(n + 2) {}
    void init(int n, COST inf = 0) {
        N = n, S = n, T = n + 1;
        OFFSET = 0, INF = inf;
        list.assign(N + 2, Edge());
        pos.clear();
    }
    friend ostream& operator << (ostream& s, const TwoVariableSubmodularOpt &G) {
        const auto &edges = G.get_edges();
        for (const auto &e : edges) s << e << endl;
        return s;
    }

    // add 1-Variable submodular functioin
    void add_single_cost(int xi, COST false_cost, COST true_cost) {
        assert(0 <= xi && xi < N);
        if (false_cost >= true_cost) {
            OFFSET += true_cost;
            add_edge(S, xi, false_cost - true_cost);
        } else {
            OFFSET += false_cost;
            add_edge(xi, T, true_cost - false_cost);
        }
    }
    
    // add "project selection" constraint (xi = T, xj = F is penalty C)
    void add_psp_penalty(int xi, int xj, COST cost) {
        assert(0 <= xi && xi < N);
        assert(0 <= xj && xj < N);
        assert(cost >= 0);
        add_edge(xi, xj, cost);
    }
    
    // add general 2-variable submodular function
    // (xi, xj) = (F, F): A, (F, T): B
    // (xi, xj) = (T, F): C, (T, T): D
    void add_submodular_function(int xi, int xj, COST A, COST B, COST C, COST D) {
        assert(0 <= xi && xi < N);
        assert(0 <= xj && xj < N);
        assert(B + C >= A + D);  // assure submodular function
        OFFSET += A;
        add_single_cost(xi, 0, D - B);
        add_single_cost(xj, 0, B - A);
        add_psp_penalty(xi, xj, B + C - A - D);
    }
    
    // add all True profit
    void add_all_true_profit(const vector<int> &xs, COST profit) {
        assert(profit >= 0);
        int slack = (int)list.size();
        list.resize(slack + 1);
        OFFSET -= profit;
        add_edge(S, slack, profit);
        for (auto xi : xs) {
            assert(xi >= 0 && xi < N);
            add_edge(slack, xi, INF);
        }
    }
    
    // solve
    COST solve() {
        return dinic() + OFFSET;
    }
    vector<bool> reconstruct() {
        vector<bool> res(N);
        queue<int> que;
        que.push(S);
        while (!que.empty()) {
            int v = que.front();
            que.pop();
            if (S < N) res[S] = true;
            for (const auto &e : list[v]) {
                if (e.cap && !res[e.to]) {
                    res[e.to] = true;
                    que.push(e.to);
                }
            }
        }
        return res;
    }
    
private:
    // add edge
    Edge &get_rev_edge(const Edge &e) {
        if (e.from != e.to) return list[e.to][e.rev];
        else return list[e.to][e.rev + 1];
    }
    Edge &get_edge(int i) {
        return list[pos[i].first][pos[i].second];
    }
    const Edge &get_edge(int i) const {
        return list[pos[i].first][pos[i].second];
    }
    vector<Edge> get_edges() const {
        vector<Edge> edges;
        for (int i = 0; i < (int)pos.size(); ++i) {
            edges.push_back(get_edge(i));
        }
        return edges;
    }
    void add_edge(int from, int to, COST cap) {
        if (!cap) return;
        pos.emplace_back(from, (int)list[from].size());
        list[from].push_back(Edge((int)list[to].size(), from, to, cap));
        list[to].push_back(Edge((int)list[from].size() - 1, to, from, 0));
    }
    
    // Dinic's algorithm
    COST dinic(COST limit_flow) {
        COST current_flow = 0;
        vector<int> level((int)list.size(), -1), iter((int)list.size(), 0);
        
        // Dinic BFS
        auto bfs = [&]() -> void {
            level.assign((int)list.size(), -1);
            level[S] = 0;
            queue<int> que;
            que.push(S);
            while (!que.empty()) {
                int v = que.front();
                que.pop();
                for (const Edge &e : list[v]) {
                    if (level[e.to] < 0 && e.cap > 0) {
                        level[e.to] = level[v] + 1;
                        if (e.to == T) return;
                        que.push(e.to);
                    }
                }
            }
        };
        
        // Dinic DFS
        auto dfs = [&](auto self, int v, COST up_flow) {
            if (v == T) return up_flow;
            COST res_flow = 0;
            for (int &i = iter[v]; i < (int)list[v].size(); ++i) {
                Edge &e = list[v][i], &re = get_rev_edge(e);
                if (level[v] >= level[e.to] || e.cap == 0) continue;
                COST flow = self(self, e.to, min(up_flow - res_flow, e.cap));
                if (flow <= 0) continue;
                res_flow += flow;
                e.cap -= flow, e.flow += flow;
                re.cap += flow, re.flow -= flow;
                if (res_flow == up_flow) break;
            }
            return res_flow;
        };
        
        // flow
        while (current_flow < limit_flow) {
            bfs();
            if (level[T] < 0) break;
            iter.assign((int)iter.size(), 0);
            while (current_flow < limit_flow) {
                COST flow = dfs(dfs, S, limit_flow - current_flow);
                if (!flow) break;
                current_flow += flow;
            }
        }
        return current_flow;
    };
    COST dinic() {
        return dinic(numeric_limits<COST>::max());
    }
};


// ABC 326 G - Unlock Achievement
void ABC_326_G() {
    int N, M;
    cin >> N >> M;
    vector<long long> C(N), A(M);
    vector<vector<long long>> L(M, vector<long long>(N));
    for (int i = 0; i < N; ++i) cin >> C[i];
    for (int i = 0; i < M; ++i) cin >> A[i];
    for (int i = 0; i < M; ++i) for (int j = 0; j < N; ++j) cin >> L[i][j];
    
    const long long INF = 1LL<<55;
    TwoVariableSubmodularOpt<long long> tvs(N*4, INF);
    for (int i = 0; i < N*4; ++i) {
        tvs.add_single_cost(i, 0, C[i/4]);
        if (i % 4 != 3) tvs.add_psp_penalty(i+1, i, INF);
    }
    for (int i = 0; i < M; ++i) {
        vector<int> ids;
        for (int j = 0; j < N; ++j) {
            if (L[i][j] > 1) ids.push_back(j*4 + (L[i][j] - 2));
        }
        tvs.add_all_true_profit(ids, A[i]);
    }
    cout << -tvs.solve() << endl;
}

int main() {
    ABC_326_G();
}