Codeforces 1129D. Isolation (2900)

给出一个长度为 \(n\) 的序列 \(a\)，把它划分成若干段，使得每一段中出现过恰好一次的元素个数 \(\le k\)，求方案数对 \(998244353\) 取模后的结果。
\(1\le k\le n\le 10^5,1\le a_i\le n\)。

---

考虑优化一个 trivial 的 \(\text{dp}\) 式子：

\[dp_i=\sum_{j} dp_{j}\cdot [cnt[j+1\ldots i]\le k] \]

容易注意到 \(cnt[j+1\ldots i]\) 每次在 \(i\) 变化到 \(i+1\) 时的变化量只为 \(\pm1\)。

• \(\forall j\in[pre_{pre_{i}},pre_{i}-1],cnt[j+1\ldots i]=cnt[j+1\ldots i-1]+1\)。
• \(\forall j\in[pre_{i},i-1],cnt[j+1\ldots i]=cnt[j+1\ldots i-1]-1\)。

所以实现的数据结构要支持区间键值 \(+1/-1\)，求区间键值 \(\le k\) 的权值和。

考虑使用分块来维护。由于每次变化量为 \(\pm 1\)，直接维护每个块键值 \(\le k\) 的权值前缀和。此时散块中的每个加一和每个减一都可以 \(O(1)\) 处理，整块直接打标记。因此修改的复杂度为 \(O(\sqrt{n})\)。

但是在对 \(i\) 求完 \(\text{dp}\) 值，还要把它添加进当前块的前缀和中，如果直接扫一遍，复杂度就变成 \(O(n)\)。处理的办法是求键值 \(\ge k\) 的权值前缀和以及每个块的所有的 \(\text{dp}\) 值的和，这样就可以 \(O(1)\) 添加当前 \(\text{dp}\) 值。

总时间复杂度 \(O(n\sqrt n)\)。

\(\color{blue}{\text{code}}\)

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+5,mod=998244353;
int n,k,B,cnt[N],dp[N],sum[330][N],bel[N],tag[N],pre[N],lst[N],tot[N],L[N],R[N],a[N];
inline void add(int l,int r,int delta){
	if(l>r)return;
	if(bel[l]==bel[r]){
		for(int i=l;i<=r;++i)
			if(delta==1)(sum[bel[l]][++cnt[i]]+=dp[i])%=mod;
			else (sum[bel[l]][cnt[i]--]+=mod-dp[i])%=mod;
		return;
	}
	for(int i=l;i<=R[bel[l]];++i)
		if(delta==1)(sum[bel[l]][++cnt[i]]+=dp[i])%=mod;
		else (sum[bel[l]][cnt[i]--]+=mod-dp[i])%=mod;
	for(int i=L[bel[r]];i<=r;++i)
		if(delta==1)(sum[bel[r]][++cnt[i]]+=dp[i])%=mod;
		else (sum[bel[r]][cnt[i]--]+=mod-dp[i])%=mod;
	for(int i=bel[l]+1;i<bel[r];++i)tag[i]+=delta;
}
inline int calc(){
	int ans=0;
	for(int i=0;i<=B;++i)
		if(k-tag[i]+1>=0)(ans+=(tot[i]-sum[i][k-tag[i]+1]+mod)%mod)%=mod;
	return ans;
}
int main(){
	scanf("%d%d",&n,&k),B=sqrt(n)+1,dp[0]=sum[1][0]=tot[1]=1;
	for(int i=1;i<=n;++i)scanf("%d",a+i),pre[i]=lst[a[i]],lst[a[i]]=i;
	for(int i=1;i<=B;++i){
		L[i]=(i-1)*B,R[i]=min(i*B-1,n);
		for(int j=L[i];j<=R[i];++j)bel[j]=i;
	}
	for(int i=1;i<=n;++i){
		add(pre[pre[i]],pre[i]-1,-1),add(pre[i],i-1,1),dp[i]=calc();
		(sum[bel[i]][0]+=dp[i])%=mod,(tot[bel[i]]+=dp[i])%=mod;
	}
	return printf("%d\n",dp[n]),0;
}