どこかでめっちゃ似たのを解いたことあると思ったらコレだった!
問題概要
長さ の正の整数列 と、正の整数 が与えられる。
個の整数 の空でない部分集合 は 通りあるが、そのそれぞれについて、
- = の部分集合のうち、総和が になるようなものの個数
として定義する。 通りの部分集合 に対する の値の総和を、 で割ったあまりを求めよ。
制約
解法 (1): 自分の考えたこと
この問題の簡単バージョンとして、 が全体集合の場合のみを考える、というのがある。その場合は、単純に
- dp[ i ][ j ] := 最初の i 個の数の部分集合であって、総和が j になるものの個数
として、
- dp[ i + 1 ][ j ] = dp[ i ][ j ] + dp[ i ][ j - A[i] ] (j - A[i] < 0 のときは前半のみ)
とかで OK。とりあえずこの手の問題の定石として、次のように考えてみる。
- 各 T に対して
- 総和が j になるような T の部分集合が何個あるかを求めて
- それを合算する
を求めるという風に見るのではなく、
- 総和が j になるようなそれぞれの部分集合 に対して
- を部分集合として含むような が何個あるかを求めて
- のサイズを とすると 通り
- それを合算する
という風に捉え直すことにする。こういう風な考え方を個人的には「横に見るものを縦に見る」と呼んでいる。その考え方を使うような問題にはタグ付けしている:
https://drken1215.hatenablog.com/archive/category/%E7%B8%A6%E3%81%AB%E8%A6%8B%E3%82%8B%E3%82%82%E3%81%AE%E3%82%92%E6%A8%AA%E3%81%AB%E8%A6%8B%E3%82%8Bdrken1215.hatenablog.com
この視点で見ると、とりあえず の解法が得られる。
- dp[ i ][ j ][ k ] := 最初の i 個の数の部分集合であって、そのうちの k 個を選んで、総和が j であるようなものの個数
これは上と同じような DP で求められる。そして
- dp[ i ][ j ][ k ] ×
を合計すれば OK。
高速化
よくよく考えると、単純に
- dp[ i ][ j ] := 最初の i 個の数のみについて考えたときの、総和が j であるように選んでできる部分集合すべてについての、それを含むような集合の個数の総和
としても遷移が作れることがわかる。i から i+1 への遷移は、次のようできる。
A[i] を選ばないときは、総和が j になるような部分集合自体は変化しないが、それを含む集合の個数はむしろ 2 倍になる (A[i] を「入れる」「入れない」)。よって、dp[ i + 1 ][ j ] += dp[ i ][ j ] × 2
A[i] を選ぶときは、dp[ i + 1 ][ j ] += dp[ i ][ j - A[i] ]
これなら計算量は になって間に合う。
#include <bits/stdc++.h> using namespace std; const int MOD = 998244353; int main() { int N, S; cin >> N >> S; vector<int> A(N); for (int i = 0; i < N; ++i) cin >> A[i]; vector<vector<long long>> dp(N+1, vector<long long>(S+1, 0)); dp[0][0] = 1; for (int i = 0; i < N; ++i) { for (int j = 0; j <= S; ++j) { dp[i+1][j] += dp[i][j] * 2 % MOD; if (j-A[i] >= 0) dp[i+1][j] += dp[i][j-A[i]]; dp[i+1][j] %= MOD; } } cout << dp[N][S] << endl; }
解法 (2): 形式的冪級数からの導出
結構この方針で解法を導いた方は多かったみたい。maspy さんの記事がとてもわかりやすい!
関連ツイート
各部分集合についてその部分集合を全探索するのはO(4^N)ではなくO(3^N)ですよという話があって、F問題はそういう思想が含まれていたりします
— てんぷら (@tempura_cpp) 2020年5月31日
F: (2+A[1]X)(2+A[2]X)...(2+A[N]X) を展開したときの X^S の係数です(この手法、ゆきこ → NOMURA → ABC で3日連続です)。
— きり (@kiri8128) 2020年5月31日
Python だと多倍長でやると簡単かつ高速です。
O(N+SlogS)に訂正します
— hotman (@hotmanww) 2020年5月31日
[x^S]Π(2+A_ix)
— hotman (@hotmanww) 2020年5月31日
=2^N*[x^S]Π(1+(A_i/2)x)
=2^N*[x^S]exp(∑log(1+(A_i/2)x))
からマクローリン展開で愚直計算+expをとる
です