久しぶりに bit ベクター高速化を使った。デバッグがしんどかった。
問題概要
長さ の数列 と、長さ の数列 が与えられる。今、次の操作をちょうど 回実行する
- と を選ぶ
- と とを swap する
操作実行後の の値の最大値を求めよ。
制約
考えたこと
どのように操作しても、 の総和と の総和の和は一定 ( とおく) なので、 を最大にするためには、 と の差を最小にしたい。
言い換えれば、操作後の の値をできるだけ に近づけたい。部分和問題っぽい雰囲気なので DP で解けそうだ。いつもの DP によって、
adp[k][v]
← 数列 から 個選んで総和を にできるかどうかbdp[k][v]
← 数列 から 個選んで総和を にできるかどうか
が求められる。これらを求めるのに計算量は を要するので通常は間に合わない。そこで常套手段として、bitset を用いた高速化をする。
集計
基本的には adp[i][x]
が True となる x と bdp[N-i][y]
が True となる y の和 x + y が に近づくように集計する。
ここで、
- のときは、 であることが条件
- のときは、 であることが条件
ということに注意しておく。
上記の条件を満たす各 に対して、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)
の最大値を求めていく。この部分の計算量は となる。十分間に合う。
コード
#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; }