kmjp's blog

競技プログラミング参加記です

Hello 2024 : E. Counting Prefixes

まぁまぁの時間で解けている。
https://codeforces.com/contest/1919/problem/E

問題

1,-1で構成されるN要素の数列Aに対し、prefix sumを取ってソートした数列Pが与えられる。
Pを構築できるAは何通りあるか。

解法

ソート前のPは、隣接要素の差は1か-1である。
dp(v,n) := Pのうち小さい順にvまで並びを決めたとき、n個の非連結な部分列を構成するような組み合わせ
として、vの小さい順にPの値の埋め方を決めていく。
新しい部分列を増やしたり減らしたりしながら定めていこう。

int T,N,P[5050];
const ll mo=998244353;
ll from[2][2][5050];
ll to[2][2][5050];

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}
ll hcomb(ll P_,ll Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		map<int,int> M;
		M[0]++;
		FOR(i,N) {
			cin>>P[i];
			M[P[i]]++;
		}
		P[N]=0;
		sort(P,P+N);
		FOR(i,N-1) if(P[i]+1<P[i+1]) break;
		if(i!=N-1) {
			cout<<0<<endl;
			continue;
		}
		sort(P,P+N+1);
		FOR(i,N) if(P[i]+1<P[i+1]) break;
		if(i!=N) {
			cout<<0<<endl;
			continue;
		}
		ZERO(from);
		from[1][1][M[P[0]]]=1;
		for(i=P[0]+1;i<=P[N];i++) {
			int a=M[i-1];
			int b=M[i];
			ZERO(to);
			FOR(x,2) FOR(y,2) FOR(k,a+1) if(from[x][y][k]) {
				ll v=from[x][y][k];
				// ->00
				if(b>=k-1) (to[0][0][1+(b-(k-1))]+=v*hcomb(k-1,b-(k-1)))%=mo;
				if(x&&b>=k) (to[1][0][1+(b-k)]+=v*hcomb(k,b-k))%=mo;
				if(y&&b>=k) (to[0][1][1+(b-k)]+=v*hcomb(k,b-k))%=mo;
				if(x&&y&&b>=k+1) (to[1][1][1+(b-(k+1))]+=v*hcomb(k+1,b-(k+1)))%=mo;
			}
			
			
			swap(from,to);
			if(i==0) {
				FOR(j,b+1) from[0][0][j]=from[0][1][j]=0;
			}
			if(i>0) {
				FOR(j,b+1) from[1][0][j]=from[1][1][j]=0;
			}
		}
		ll ret=0;
		ret+=from[0][0][1]+from[0][1][1]+from[1][0][1]+from[1][1][1];
		cout<<ret%mo<<endl;
		
	}
}

まとめ

以前ARCで似たような問題を解けなかったことを覚えていて、おかげですんなり回答にたどり着けた。