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

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

DISCO presents 2016 予選 D - DDPC特別ビュッフェ (赤色)

久しぶりに bit ベクター高速化を使った。デバッグがしんどかった。

問題概要

長さ  N の数列  A_{1}, A_{2}, \dots, A_{N} と、長さ  M の数列  B_{1}, B_{2}, \dots, B_{M} が与えられる。今、次の操作をちょうど  K 回実行する

  •  1 \le i \le N 1 \le j \le M を選ぶ
  •  A_{i} B_{j} とを swap する

操作実行後の  \displaystyle \bigl(\sum_{i=1}^{N} A_{i}\bigr)\bigl(\sum_{i=1}^{M} B_{i}\bigr) の値の最大値を求めよ。

制約

  •  1 \le N, M \le 55
  •  1 \le K \le 999
  •  0 \le A_{i}, B_{i} \le 22222

考えたこと

どのように操作しても、 A の総和と  B の総和の和は一定 ( S とおく) なので、 \displaystyle \bigl(\sum_{i=1}^{N} A_{i}\bigr)\bigl(\sum_{i=1}^{M} B_{i}\bigr) を最大にするためには、 \displaystyle \bigl(\sum_{i=1}^{N} A_{i}\bigr) \displaystyle \bigl(\sum_{i=1}^{N} B_{i}\bigr) の差を最小にしたい。

言い換えれば、操作後の  \displaystyle \bigl(\sum_{i=1}^{N} A_{i}\bigr) の値をできるだけ  \frac{S}{2} に近づけたい。部分和問題っぽい雰囲気なので DP で解けそうだ。いつもの DP によって、

  • adp[k][v] ← 数列  A から  k 個選んで総和を  v にできるかどうか
  • bdp[k][v] ← 数列  B から  k 個選んで総和を  v にできるかどうか

が求められる。これらを求めるのに計算量は  O((N^{2}+M^{2})S) を要するので通常は間に合わない。そこで常套手段として、bitset を用いた高速化をする。

集計

基本的には adp[i][x] が True となる x と bdp[N-i][y] が True となる y の和 x + y が  \frac{S}{2} に近づくように集計する。

ここで、

  •  K = 1 のときは、 i = N-1, N-i = 1 であることが条件
  •  K > 1 のときは、 N-i \le \min(K, M) であることが条件

ということに注意しておく。

上記の条件を満たす各  i に対して、adp[i][x] == True であるような x に対して、

  • S/2 以下の y であって bdp[N-i][y] == Trueとなる最大の y
  • (S+1)/2 以上の y であって bdp[N-i][y] == Trueとなる最小の y

をそれぞれ求めて、(x+y)(S-x-y) の最大値を求めていく。この部分の計算量は  O(NS) となる。十分間に合う。

コード

#include <bits/stdc++.h>
using namespace std;
constexpr long long MAX = 1500000;

// 前処理の DP
vector<bitset<MAX>> enum_sum(const vector<int>& A) {
    int N = A.size();
    vector<bitset<MAX>> dp(N+1);
    dp[0][0] = true;
    for (int i = 0; i < N; ++i) {
        vector<bitset<MAX>> ndp(N+1);
        for (int j = 0; j <= N; ++j) {
            // 選ぶ
            if (j < N) ndp[j+1] |= dp[j] << A[i];

            // 選ばない
            ndp[j] |= dp[j];
        }
        swap(dp, ndp);
    }
    return dp;
}

// 集合 a と集合 b から要素を選んで和を S/2 に近づけたい
long long solve(const bitset<MAX>& a, const bitset<MAX>& b, long long S) {
    long long res = 0;

    // left[v] := (x <= v かつ b[x] = true となる最大の x)
    // right[v] := (x >= v かつ b[x] = true となる最小の x)
    vector<long long> left(MAX, -1), right(MAX, -1);
    for (int i = 0; i < MAX; ++i) {
        if (b[i]) left[i] = i;
        else if (i-1 >= 0) left[i] = left[i-1];
    }
    for (int i = MAX-1; i >= 0; --i) {
        if (b[i]) right[i] = i;
        else if (i+1 < MAX) right[i] = right[i+1];
    }

    // 集合 a の値 v に応じて、最適な b を選ぶ
    for (int v = 0; v < MAX; ++v) {
        if (!a[v]) continue;
        long long bv1 = min(max(S/2-v, 0LL), MAX);
        long long bv2 = min(max((S+1)/2-v, 0LL), MAX);
        if (left[bv1] != -1)
            res = max(res, (v + left[bv1]) * (S - v - left[bv1]));
        if (right[bv2] != -1)
            res = max(res, (v + right[bv2]) * (S - v - right[bv2]));
    }
    return res;
}

int main() {
    // 入力
    int N, M, K;
    cin >> N >> M >> K;
    vector<int> A(N), B(M);
    long long S = 0;
    for (int i = 0; i < N; ++i) { cin >> A[i]; S += A[i]; }
    for (int i = 0; i < M; ++i) { cin >> B[i]; S += B[i]; }

    // adp[i] := A から i 個選んで作れる数を表す bitset
    // bdp[j] := B から j 個選んで作れる数を表す bitset
    const auto& adp = enum_sum(A);
    const auto& bdp = enum_sum(B);

    // adp[i] と bdp[N-i] から S/2 に近い数を作る
    long long res = 0;
    for (int i = 0; i <= N; ++i) {
        if (K == 1 && i == N) continue;
        if (N-i > min(K, M)) continue;
        res = max(res, solve(adp[i], bdp[N-i], S));     
    }
    cout << res << endl;
}