2023-09-25

AtCoder ABC 321 F - #(subset sum = K) with Add and Erase (青色, 525 点)

制約条件:総和=K AtCoder 多項式・FPS(形式的冪級数) FPSテク:漸化式の母関数を求める戻すDP AtCoder525点 ABC-F 青色diff DP シーケンシャルDP ナップサック操作:削除操作:挿入操作部分和各kに対して

「戻す DP」または「FPS」で！

問題へのリンク

問題概要

最初、箱は空である。以下の操作を $Q$ 回行う。各操作後において、以下の値を答えよ。

箱に入っているボールをいくつか選ぶ方法であって、ボールに書かれた数値の総和が $K$ となる方法の個数を 998244353 で割った余り

各操作は次の 2 タイプある。

+ v：整数値 $v$ の書かれたボールを 1 つ箱に追加する
- v：整数値 $v$ の書かれたボールを 1 つ箱から取り出す (1 個以上存在することが保証される)

制約

$1 \le Q, K \le 5000$

解法 (1)：戻す DP

部分和問題 (数え上げ ver.) の復習

今回の問題は、もし最後の状態の箱について問題を解くだけなら、いつもの DP でできる。部分和問題を数え上げ問題にしたような問題となる。

ここで、一旦、部分和問題 (数え上げ ver.) に対する DP を復習してみよう。

$N$ 個の要素 $v_{0}, v_{1}, \dots, v_{N-1}$ からいくつか選ぶ方法のうち、総和が $K$ となるものの個数を求めよ

この問題は次のような DP で解ける。

dp[i][j] ← $N$ 個の要素のうち最初の $i$ 個の要素のみからいくつか選ぶ方法のうち、総和が $j$ となるものの個数

そして、次の更新式を組める。

dp[i+1][j] = dp[i][j] + (j >= v[i] ? dp[i][j-v[i]] : 0)

この更新式は見方を変えると、サイズ $K+1$ の配列 dp[i] に対して、要素 $v$ を用いて、サイズ $K+1$ の配列 dp[i+1] へと変換したことを表している。

この変換は、 $K+1$ 次元ベクトルから $K+1$ 次元ベクトルへの関数とみなすこともできる。これを具体的に表現してみよう。

DP の更新式を関数で表す

要素 $v$ に対して、 $K+1$ 次元ベクトル $x = (x_{0}, x_{1}, \dots, x_{K})$ から $K+1$ 次元ベクトルへの関数 $f_{v}$ を次のように定義する。

$f_{v}(x_{j}) = \left\{ \begin{array}{ll} x_{j} + x_{j-v} & (j \ge v)\\ x_{j} & (j \lt v) \end{array} \right.$

このとき、DP の更新式は次のように書ける。

dp[i+1] = $f_{v}$ (dp[i])

戻す DP

いよいよ、「戻す DP」の核心に迫る。一般に、

$m-1$ 個の要素 $v_{0}, v_{1}, \dots, v_{m-2}$ からいくつかを選んで総和を $j$ にする方法を $x_{j}$ と表す
$m$ 個の要素 $v_{0}, v_{1}, \dots, v_{m-2}, v$ からいくつかを選んで総和を $j$ にする方法を $y_{j}$ と表す

ことにする。 $x, y$ は $K+1$ 次元ベクトルとなる。このとき、

$y = f_{v}(x)$

と表せる。ここで関数 $f_{v}$ の逆関数 $f_{v}^{-1}$ を求めることができれば、

$x = f_{v}^{-1}(y)$

と表すことができる。これが戻す DP だ。つまり、「 $m$ 個の要素 $v_{0}, v_{1}, \dots, v_{m-2}, v$ に関する部分和問題の結果」から、「 $m-1$ 個の要素 $v_{0}, v_{1}, \dots, v_{m-2}$ に関する部分和問題の結果」を得ることができるのだ。

逆関数 $f_{v}^{-1}$ を求める

よって、残りは「関数 $f_{v}$ の逆変換 $f_{v}^{-1}$ を求める」ことができればよい。高校数学でもやった通り、 $x, y$ を入れ替えてみる。

$y_{j} = \left\{ \begin{array}{ll} x_{j} + x_{j-v} & (j \ge v)\\ x_{j} & (j \lt v) \end{array} \right.$

において $x, y$ を入れ替えると、

$x_{j} = y_{j} + y_{j-v}$ ( $j \ge v$ )
$x_{j} = y_{j}$ ( $j \lt v$ )

となる。 $y$ を主役にすると

$y_{j} = x_{j} - y_{j-v}$ ( $j \ge v$ )
$y_{j} = x_{j}$ ( $j \lt v$ )

となる。これは、for 文で実装すると、次のようになる。

for (int j = 0; j <= K; ++j) {
    y[j] = x[j];
    if (j >= v) y[j] -= y[j - v];
}

実装上の整理

ここまでの結果をまとめよう。

配列 dp に対して、要素 $v$ を追加して、配列 dp2 を得る操作は次のように書ける (普通の部分和 DP)。

for (int j = 0; j <= K; ++j) {
    dp2[j] = dp[j];
    if (j >= v) dp2[j] += dp[j - v];
}

そして、配列 dp に対して、要素 $v$ を削除して、配列 dp2 を得る操作は次のように書ける (戻す DP)。

for (int j = 0; j <= K; ++j) {
    dp2[j] = dp[j];
    if (j >= v) dp2[j] -= dp2[j - v];
}

解答例コード

これらをまとめて、次のコードで AC できる。計算量は $O(QK)$ となる。

#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 を直接弄る実装もできる。次のように書ける。とても簡潔になると同時に、「 $v$ を追加」「 $v$ を削除」が互いに逆変換であることがハッキリ見て取れるようになる。

#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 で考えると意味合いがより明確になる。それについては、公式解説を参照。

atcoder.jp

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

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

AtCoder ABC 321 F - #(subset sum = K) with Add and Erase (青色, 525 点)

問題概要

制約

解法 (1)：戻す DP

部分和問題 (数え上げ ver.) の復習

DP の更新式を関数で表す

戻す DP

逆関数 $f_{v}^{-1}$ を求める

実装上の整理

解答例コード

実装を洗練化：in-place 化

解法 (2)：FPS

問題概要

制約

解法 (1)：戻す DP

部分和問題 (数え上げ ver.) の復習

DP の更新式を関数で表す

戻す DP

逆関数 を求める

実装上の整理

解答例コード

実装を洗練化：in-place 化

解法 (2)：FPS

逆関数 $f_{v}^{-1}$ を求める