AtCoder Regular Contest 060 C - 高橋君とカード

問題

atcoder.jp

問題概要

N枚のカードがあり、i番目(1-index)のカードには整数x[i-1]が書かれている。
何枚かカードを選びカードの整数の平均がAとなるようにしたい。そのような選び方は何通りあるか。

解法

平均を平均のまま扱うのは大変なので合計と個数を扱う。

dp[i][j][k] : i番目(1-index)まで見てj枚選び、合計がkである通り数。
0≦i≦N、0≦j≦N、0≦k≦Σx[i-1]

こうすれば答えは
i=Nかつk=j*Aとなるdp[i][j][k]の合計である。(j=0かつk=0の時を除く)

i番目のカードを選ぶか選ばないかの2通りある。

i番目のカードを選ばないとき
dp[i][j][k]=dp[i-1][j][k]

i番目のカードを選ぶとき
dp[i][j][k]=dp[i-1][j-1][k-x[i-1]]

これらの合計となる。
しかしj-1<0またはk-x[i-1]<0の時は配列外参照となるため
dp[i][j][k]=dp[i-1][j][k]
となる。

i=0の時
j=0かつk=0の時はdp[i][j][k]=1
それ以外の時はdp[i][j][k]=0

ソースコード

signed main() {
	int N, A;
	cin >> N >> A;
	vector<int>x(N);
	int Sx = 0;
	rep(i, N) {
		cin >> x[i];
		Sx += x[i];
	}
	vector<vector<vector<int>>>dp(N + 1, vector<vector<int>>(N + 1, vector<int>(Sx + 1)));
	rep(i, N + 1) {
		rep(j, N + 1) {
			rep(k, Sx + 1) {
				if (i == 0) {
					if (j == 0 && k == 0) {
						dp[i][j][k] = 1;
					}
					else {
						dp[i][j][k] = 0;
					}
				}
				else {
					if (j - 1 < 0 || k - x[i - 1] < 0) {
						dp[i][j][k] = dp[i - 1][j][k];
					}
					else {
						dp[i][j][k] = dp[i - 1][j][k] + dp[i - 1][j - 1][k - x[i - 1]];
					}
				}
			}
		}
	}
	int ans = 0;
	rep(i, N + 1) {
		rep(j, N + 1) {
			rep(k, Sx + 1) {
				if (i == N && k == j * A && j!=0 && k!=0) {
					ans += dp[i][j][k];
				}
			}
		}
	}
	cout << ans << endl;
}