Codeforces #542 Div1 D. Isolation

いやこれありなのか…。
https://codeforces.com/contest/1129/problem/D

問題

N要素の整数列A[i]が与えられる。また、パラメータKが与えられる。
整数列を以下の条件を満たしながら分割したい。

• 分割した各整数列において、1回しか登場しない要素はK個以下である。

分割の仕方は何通りか。

解法

これ、実はGCCの最適化オプション次第でO(N^2)解法が通ってしまう。
f(n) := A[0...(n-1)]における題意を満たす分割の個数
とする。ここにA[n]を加えf(n+1)を求めることを考えよう。

まず数列の各要素iに対し、A[i..n]において1回しか登場しない要素の数をg(i)とする。
また、ある値について、直前に登場した位置Pと2つ前に登場した位置Qを考える。

A[n]が末尾に加わることで、g(P+1)～g(Q)は1減少し、g(Q+1)～g(n)は1加算する。
f(n+1)=f(n)としたうえで、上記g(*)の変化によってKを超過、ないし下回った位置kに対し、f(n+1)+=f(k)と足しこんでいけばよい。

ただ、これはO(N^2)解法なので想定解法ではない。
想定解法は平方分割っぽいね。

#pragma GCC optimize("Ofast")
#pragma GCC target("avx,avx2")

int N,K;
int A[101010];
int B1[101010],B2[101010];
int num[101010];
int dp[101010];
ll ret;
int mo=998244353;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	dp[0]=1;
	ll sum=1;
	for(i=1;i<=N;i++) {
		cin>>x;
		for(j=B2[x];j<B1[x];j++) if(--num[j]==K) sum+=dp[j];
		for(j=B1[x];j<i;j++) if(++num[j]==K+1) sum-=dp[j];
		
		sum%=mo;
		dp[i]=sum;
		B2[x]=B1[x];
		B1[x]=i;
		
		sum=(sum+dp[i])%mo;
		if(sum>=mo) sum-=mo;
	}
	cout<<(dp[N]%mo+mo)%mo<<endl;
}

解法

これ32bit環境だからなおさら効果が高いってのもありそう。