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

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

AtCoder ARC 060 C - 高橋君とカード (300 点)

「平均値が A」という制約は上手に扱うテクニックがある。
今回はそれを思いつかなくても解けるけど、思いつくと計算量が落ちる。

問題へのリンク

問題概要

 N 個の整数  x_{1}, \dots, x_{N} からいくつか選ぶ方法のうち、その平均値がちょうど  A となるものが何通りあるかを求めよ。

制約

  •  1 \le N, A \le 50

解法 (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 ]

で表すことができる。計算量は  O(N^{3}A)

#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): 平均値制約を上手に

一般に

 a, b, c, \dots の平均値が  A
 a-A, b-A, c-A, \dots の平均値が  0
 a-A, b-A, c-A, \dots の総和が  0

という言い換えができる。平均値に関する制約が、総和に関する制約になって、とても扱いやすくなる!!!元の問題は


 N 個の整数  x_{1} - A, \dots, x_{N} - A が与えられる。これらからいくつか選んで総和を 0 にする方法が何通りあるかを求めよ。


と言い換えることができる。こうしておくと dp も、

dp[ i ][ s ] := 最初の i 個からいくつか選んで総和が s になる方法の個数

という風に次元を減らすことができる。計算量は  O(N^{2}A) に減る。

#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;
}