高橋君とカード #C++

高橋君とカード / Tak and Cards

考察

DFSで全探索すると$O(2^N)$で間に合わない
貪欲だと厳しいかな…?普通に無理そう
とりあえず全探索を書いてみる
最終的には平均値を求めるわけだが、これがちょっと複雑。小数点とか出てきそうだし。最終的に選んだカードの枚数を$m$枚とすると、平均の式は$\frac{m枚の合計}{m枚} = A$とすることができる。これを式変形すると、$m枚の合計 = A\times m枚$にできる。つまり、最終的に選んだカードの枚数とその合計がわかればそれがAになるかがわかる（後で気づいたけど、式変形とかなくても解けるな。この問題）
これまでの情報で、rec(インデックス, 選んだ枚数, 合計) = 総数って感じの関数でDFSを実装すればいいことがわかる。あとは、遷移を考えて頑張って実装する。
実装したら、正しく動くか確認する。サンプルと合わせた結果、出力された数値は答えより1大きくなってる。答えから1を引いて提出してもいいけど、なんか気持ち悪いのでデバッグしてみる。すると、選んだ枚数が0枚の時も平均値がAであるという結果を返すコードになっていた。なので、選んだ枚数が0枚の時の処理も書く。
コードをみると、メモ化再帰にできそうなのでそうする。
終了

コード

暇なのでループDPでも実装してみた。

#include <bits/stdc++.h>
using namespace std;

#define int long long

int N, A;
vector<int> x;
int dp[55][55][2555]; // dp[インデックス][枚数][合計] = 総数

int rec(int idx, int mai, int sum) {
  // 全部調べ終わる
  if (idx >= N) {
    return ((mai * A) == sum) && (sum != 0);
  }

  // 計算済み
  if (dp[idx][mai][sum] != -1) {
    return dp[idx][mai][sum];
  }

  int ret = 0;

  // index枚目のカードを選ぶ
  int yes = rec(idx + 1, mai + 1, sum + x[idx]);
  // index枚目のカードを選ばない
  int no = rec(idx + 1, mai, sum);

  ret += yes + no;

  return dp[idx][mai][sum] = ret;
}

signed main() {
  cin >> N >> A;
  x.resize(N);
  for (int i = 0; i < N; i++) {
    cin >> x[i];
  }

  memset(dp, -1, sizeof(dp));
  int ans = rec(0, 0, 0);

  cout << ans << endl;

  return 0;
}

もうちょい計算量を減らす

なんと解説PDFでは2次元DPで解説されてる。何それ怖い。
まあ、この考察がいつか使えるかもしれないので考えてみる
さっきの考察で、$ A\times m枚 = m枚の合計$という式を作った。これは、mの枚数とm枚の合計がわかっているときに平均がAになるかを判定するのに使った式ね。この式が正しければ平均をAにすることができる（これ大事）
ここで、選ぶと平均値がAになるようなカードの集合を$Y$とする。
そして、$ A\times m枚 = m枚の合計$と言う式をシグマの式に変形すると$\sum_{i \in Y}^{|Y|}A = \sum_{i \in Y}^{|Y|}x_i$となる（$|Y|$は集合Yの中身の個数みたいな意味で使ってる）
さらに式変形すると$\sum_{i \in Y}^{|Y|}(x_i - A) = 0$となる。これは選ぶ全カードからAを引いたときのカードの総和が0になることを意味している。先ほどと同様にこの式が正しい時、平均をAにすることができる。
この式を見ると、枚数を考えなくていいことがわかる。なので、DPの次元が1つ減って計算量も減る。

解法

あらかじめ全カードの値からAを引いておく。
DPする

コード

感想

1年くらい前にこの記事で解説を書いた問題。当時はあんまり理解してなかったけど今ではすんなり解けたので嬉しい（DPで解くことがわかってたからすんなり解けたってのもあるけど）
Aを引いておくことで合計が0になっていればOKってやつ、式変形すると正しいことはわかる。でも、直感的に理解できない。わかるようなわからないようなふわふわした感じになる。
あと、このDPの次元減らすやつを思いつける気がしない。式変形思いつくしかないのか？