久しぶりに 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;
}

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

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

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

問題概要

制約

考えたこと

集計

コード