H - Avoid K Partition 解説
by
Nyaan
説明のため、問題を次のように言い換えます。
地点 \(1, 2, \dots, N+1\) がこの順に並んでいる。
地点をいくつか選んで「区切り」を入れる。ただし、
- 地点 \(1, N+1\) には必ず区切りを入れる必要がある。
- 地点 \(i, j\) に区切りがあってその間に区切りが無いとする。このとき \(\sum_{k=i}^{j-1} A_k\) は \(K\) と等しくない必要がある。
条件を満たす区切りの入れ方は何通りか?\((\text{mod }998244353)\)
このように言い換えることで \(\mathrm{O}(N^2)\) の動的計画法を立てることができます。すなわち、
- \(dp[i]\) : 左から見て行って、最後に区切りを入れたのが地点 \(i\) である通り数
とすると、初期状態は \(dp[1] = 1\)、答えは \(dp[N + 1]\) で遷移が
\[dp[n] = \sum_{1 \leq m \lt n} \left(0\text{ if }\sum_{k=m}^{n-1} A_i = K \text{ else } dp[m] \right)\]
になることが確認できます。
この DP を高速化します。\(A\) の累積和テーブル \(B\) を \(B_i = \sum_{j = 1}^i A_j\) として定義します。すると上記の DP の遷移は
\[ \begin{aligned} dp[n] &= \sum_{1 \leq m \lt n} \left(0\text{ if } B_{n-1} - B_{m-1} = K \text{ else } dp[m] \right) \\ &= \sum_{1 \leq m \lt n, B_{m-1} \neq B_{n-1} - K} dp[m] \end{aligned} \]
と表現できます。これは言い換えると、\(dp[m]\) が \(dp[n]\) に寄与するかどうかは \(B_{m-1}\) の値によってのみ決まるということです。
そこで、DP テーブルを連想配列(C++ の std::map
) を用いて管理します。すなわち、連想配列 \(M\) を
- \(M[x]\) : \(B_{m-1} = x\) であるような \(m\) に対する \(dp[m]\) の総和
とします。また、別途 \(dp[m]\) の総和も変数 \(all\) に管理します。このようにすることで、DP の遷移は
\[dp[n] = all - M[B_{n-1} - K]\]
となるため高速に計算できるようになり、\(M\) や \(all\) の更新も同様に高速に処理できます。
以上の内容を適切に実装することで、この問題を \(\mathrm{O}(N \log N)\) あるいは \(\mathrm{O}(N)\) で解くことが出来て、これは十分高速です。
C++ による実装例は次の通りです。
#include <iostream>
#include <map>
#include <vector>
using namespace std;
#include "atcoder/modint.hpp"
using mint = atcoder::modint998244353;
int main() {
long long N, K;
cin >> N >> K;
vector<long long> A(N);
for (auto& a : A) cin >> a;
map<long long, mint> M;
M[0] = 1;
mint all = 1;
long long acc = 0;
for (int i = 0; i < N; i++) {
acc += A[i];
long long ban = acc - K;
mint cur = all - M[ban];
M[acc] += cur, all += cur;
if (i + 1 == N) cout << cur.val() << "\n";
}
}
投稿日時:
最終更新: