公式

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";
  }
}

投稿日時:
最終更新: