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

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

AtCoder ABC 169 F - Knapsack for All Subsets (青色, 600 点)

どこかでめっちゃ似たのを解いたことあると思ったらコレだった!

drken1215.hatenablog.com

問題へのリンク

問題概要

長さ  N の正の整数列  A_{1}, A_{2}, \dots, A_{N} と、正の整数  S が与えられる。
 N 個の整数  A_{1}, A_{2}, \dots, A_{N} の空でない部分集合  T 2^{N}-1 通りあるが、そのそれぞれについて、

  •  f(T) =  T の部分集合のうち、総和が  S になるようなものの個数

として定義する。 2^{N}-1 通りの部分集合  T に対する  f(T) の値の総和を、 998244353 で割ったあまりを求めよ。

制約

  •  1 \le N, A_{i}, S \le 3000

解法 (1): 自分の考えたこと

この問題の簡単バージョンとして、 T が全体集合の場合のみを考える、というのがある。その場合は、単純に

  • 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 になるようなそれぞれの部分集合  X に対して
  •  X を部分集合として含むような  T が何個あるかを求めて
    •  X のサイズを  k とすると  2^{N-k} 通り
  • それを合算する

という風に捉え直すことにする。こういう風な考え方を個人的には「横に見るものを縦に見る」と呼んでいる。その考え方を使うような問題にはタグ付けしている:

drken1215.hatenablog.com

この視点で見ると、とりあえず  O(N^{2}S) の解法が得られる。

  • dp[ i ][ j ][ k ] := 最初の i 個の数の部分集合であって、そのうちの k 個を選んで、総和が j であるようなものの個数

これは上と同じような DP で求められる。そして

  • dp[ i ][ j ][ k ] ×  2^{N-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] ]

これなら計算量は  O(NS) になって間に合う。

#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 さんの記事がとてもわかりやすい!

maspypy.com

関連ツイート