「戻す DP」または「FPS」で！
問題概要
最初、箱は空である。以下の操作を 回行う。各操作後において、以下の値を答えよ。
箱に入っているボールをいくつか選ぶ方法であって、ボールに書かれた数値の総和が となる方法の個数を 998244353 で割った余り
各操作は次の 2 タイプある。
+ v：整数値の書かれたボールを 1 つ箱に追加する
- v：整数値
制約
解法 (1)：戻す DP
部分和問題 (数え上げ ver.) の復習
今回の問題は、もし最後の状態の箱について問題を解くだけなら、いつもの DP でできる。部分和問題を数え上げ問題にしたような問題となる。
ここで、一旦、部分和問題 (数え上げ ver.) に対する DP を復習してみよう。
個の要素
からいくつか選ぶ方法のうち、総和が
となるものの個数を求めよ
この問題は次のような DP で解ける。
dp[i][j] ← 個の要素のうち最初の
個の要素のみからいくつか選ぶ方法のうち、総和が
となるものの個数
そして、次の更新式を組める。
dp[i+1][j] = dp[i][j] + (j >= v[i] ? dp[i][j-v[i]] : 0)
この更新式は見方を変えると、サイズ の配列
dp[i] に対して、要素 を用いて、サイズ
の配列
dp[i+1] へと変換したことを表している。
この変換は、 次元ベクトルから
次元ベクトルへの関数とみなすこともできる。これを具体的に表現してみよう。
DP の更新式を関数で表す
要素 に対して、
次元ベクトル
から
次元ベクトルへの関数
を次のように定義する。
このとき、DP の更新式は次のように書ける。
dp[i+1] = (
dp[i])
戻す DP
いよいよ、「戻す DP」の核心に迫る。一般に、
個の要素
からいくつかを選んで総和を
にする方法を
と表す
個の要素
からいくつかを選んで総和を
と表す
ことにする。 は
次元ベクトルとなる。このとき、
と表せる。ここで関数 の逆関数
を求めることができれば、
と表すことができる。これが戻す DP だ。つまり、「 個の要素
に関する部分和問題の結果」から、「
個の要素
に関する部分和問題の結果」を得ることができるのだ。
逆関数 
を求める
よって、残りは「関数 の逆変換
を求める」ことができればよい。高校数学でもやった通り、
を入れ替えてみる。
において を入れ替えると、
(
)
(
)
となる。 を主役にすると
(
(
となる。これは、for 文で実装すると、次のようになる。
for (int j = 0; j <= K; ++j) { y[j] = x[j]; if (j >= v) y[j] -= y[j - v]; }
実装上の整理
ここまでの結果をまとめよう。
配列 dp に対して、要素 を追加して、配列
dp2 を得る操作は次のように書ける (普通の部分和 DP)。
for (int j = 0; j <= K; ++j) { dp2[j] = dp[j]; if (j >= v) dp2[j] += dp[j - v]; }
そして、配列 dp に対して、要素 を削除して、配列
dp2 を得る操作は次のように書ける (戻す DP)。
for (int j = 0; j <= K; ++j) { dp2[j] = dp[j]; if (j >= v) dp2[j] -= dp2[j - v]; }
解答例コード
これらをまとめて、次のコードで AC できる。計算量は となる。
#include <bits/stdc++.h> #include <atcoder/modint> using namespace std; using mint = atcoder::modint998244353; int main() { int Q, K; cin >> Q >> K; vector<mint> dp(K+1, 0); dp[0] = 1; while (Q--) { char c; int v; cin >> c >> v; if (c == '+') { vector<mint> dp2(K+1, 0); for (int i = 0; i <= K; ++i) { dp2[i] = dp[i]; if (i >= v) dp2[i] += dp[i-v]; } dp = dp2; } else { vector<mint> dp2(K+1, 0); for (int i = 0; i <= K; ++i) { dp2[i] = dp[i]; if (i >= v) dp2[i] -= dp2[i-v]; } dp = dp2; } cout << dp[K].val() << endl; } }
実装を洗練化：in-place 化
上のコードでは、配列 dp を変換したものを、新たな配列 dp2 に格納して、それを dp に戻すように実装している。
しかし、配列 dp2 を用いずに配列 dp を直接弄る実装もできる。次のように書ける。とても簡潔になると同時に、「 を追加」「
を削除」が互いに逆変換であることがハッキリ見て取れるようになる。
#include <bits/stdc++.h> #include <atcoder/modint> using namespace std; using mint = atcoder::modint998244353; int main() { int Q, K; cin >> Q >> K; vector<mint> dp(K+1, 0); dp[0] = 1; while (Q--) { char c; int v; cin >> c >> v; if (c == '+') { for (int i = K; i >= v; --i) { dp[i] += dp[i-v]; } } else { for (int i = v; i <= K; ++i) { dp[i] -= dp[i-v]; } } cout << dp[K].val() << endl; } }
解法 (2)：FPS
解法 (1) で考えた DP は、FPS で考えると意味合いがより明確になる。それについては、公式解説を参照。
