「平均値が A」という制約は上手に扱うテクニックがある。
今回はそれを思いつかなくても解けるけど、思いつくと計算量が落ちる。
問題概要
個の整数 からいくつか選ぶ方法のうち、その平均値がちょうど となるものが何通りあるかを求めよ。
制約
解法 (1): 平均値をまともに取り扱う
「平均値が A」というのは
- 個数が k 個のとき、総和が Ak
という風にとらえることができる。よって
- dp[ i ][ s ][ k ] := N 個の整数のうちの最初の i 個から、k 個選ぶ場合について、総和を s にするものが何個あるか
という DP で解くことができる。遷移は
- x[ i ] を選ばない場合: dp[ i + 1 ][ s ][ k ] += dp[ i ][ s ][ k ]
- x[ i ] を選ぶ場合: dp[ i + 1 ][ s + x[ i ] ][ k + 1 ] += dp[ i ][ s ][ k ]
で表すことができる。計算量は 。
#include <bits/stdc++.h> using namespace std; long long dp[55][3000][55]; int main() { int N, A; cin >> N >> A; vector<int> x(N); for (int i = 0; i < N; ++i) cin >> x[i]; memset(dp, 0, sizeof(dp)); dp[0][0][0] = 1; for (int i = 0; i < N; ++i) { for (int s = 0; s <= N*A; ++s) { for (int k = 0; k <= N; ++k) { if (dp[i][s][k] == 0) continue; dp[i+1][s][k] += dp[i][s][k]; dp[i+1][s+x[i]][k+1] += dp[i][s][k]; } } } long long res = 0; for (int k = 1; k <= N; ++k) res += dp[N][A*k][k]; cout << res << endl; }
解法 (2): 平均値制約を上手に
一般に
の平均値が
⇔ の平均値が
⇔ の総和が
という言い換えができる。平均値に関する制約が、総和に関する制約になって、とても扱いやすくなる!!!元の問題は
個の整数 が与えられる。これらからいくつか選んで総和を 0 にする方法が何通りあるかを求めよ。
と言い換えることができる。こうしておくと dp も、
dp[ i ][ s ] := 最初の i 個からいくつか選んで総和が s になる方法の個数
という風に次元を減らすことができる。計算量は に減る。
#include <bits/stdc++.h> using namespace std; long long dp[55][5000]; const int GETA = 2500; int main() { int N, A; cin >> N >> A; vector<int> x(N); for (int i = 0; i < N; ++i) cin >> x[i], x[i] -= A; memset(dp, 0, sizeof(dp)); dp[0][GETA] = 1; for (int i = 0; i < N; ++i) { for (int s = 0; s+x[i] < 5000; ++s) { if (dp[i][s] == 0) continue; dp[i+1][s] += dp[i][s]; dp[i+1][s+x[i]] += dp[i][s]; } } // 「何も選ばない」を除く cout << dp[N][GETA] - 1 << endl; }