ABC370 E - Avoid K Partition
求一个序列的合法划分方案数。一种划分合法当且仅当没有一个子串的和是 \(k\)。
由于是否存在子串和为 \(k\) 很重要,因此考虑将它加入状态设计中,记 \(f[i][0/1]\) 表示 \(1 \sim i\),\(i\) 处结束,还没有 / 已有和为 \(k\) 的子段,方案数。
用 \(s[i]\) 表示前缀和,得到朴素的 \(O(n^2)\) 转移:
\[f[i][0] = \sum_{0 \le j \lt i \cap s[i] - s[j] \ne k} f[j][0]
\]
\[f[i][1] = \sum_{0 \le j \lt i}f[j][1] + \sum_{0 \le j \lt i \cap s[i] - s[j] = k} f[j][0]
\]
发现 \(a[i], k\) 可以是负数,因此 \(s\) 没有单调性,不好直接找每个 \(i\) 对应的 \(j\),可能也不止一个。但对于固定的 \(i\),\(s[j] = s[i] - k\) 是固定的,因此我们可以对于每个 \(s[j]\) 都实时维护 \(\sum{f[j][0]}\),值域较大,用 \(map\) 实现就 ok。最后对 \(f[][0/1]\) 分别记录前缀和就可以做到 \(O(nlogn)\)。转移没什么变化,直接看代码,最后记得取模变正数。
#include<bits/stdc++.h>
#define F(i,l,r) for(int i(l);i<=(r);++i)
#define G(i,r,l) for(int i(r);i>=(l);--i)
#define int ll
using namespace std;
using ll = long long;
const int N = 2e5 + 5;
const int mod = 998244353;
int f[N][2], s[N];
int n, k, a[N];
map<int, int> mp;
signed main(){
ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
cin >> n >> k;
F(i, 1, n) {
cin >> a[i];
s[i] = s[i - 1] + a[i];
}
int sm0 = 1, sm1 = 0;
mp[0] = 1;
F(i, 1, n){
int ret = mp[s[i] - k];
f[i][0] = (sm0 - ret) % mod;
f[i][1] = (sm1 + ret) % mod;
(mp[s[i]] += f[i][0]) %= mod;
(sm0 += f[i][0]) %= mod;
(sm1 += f[i][1]) %= mod;
}
cout << (f[n][0] + mod) % mod << '\n';
return fflush(0), 0;
}