ABC370 E - Avoid K Partition

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<i\cap s[i]-s[j]\ne k}f[j][0]$$

$$f[i][1]=\sum _{0\le j<i}f[j][1]+\sum _{0\le j<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。在 CSP-S 2024 染色 一题中，我们可以用类似技巧将 dp 优化到线性（其实就是桶）。

最后对 $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;
}