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