Pinely Round 3 : F2. Small Permutation Problem (Hard Version)

土壇場すぎる。
https://codeforces.com/contest/1909/problem/F2

問題

1～NのPermutation Pが、N要素の数列Aに対してGoodであるとは以下を意味する。

• A[i]が非負の場合、Pの先頭i要素のうちi以下のものはA[i]個である。

Aが与えられたとき、GoodとなるPは何通りか。

解法

先頭i要素目を考えるとき、A[i]が非負で、iの手前のA[j]が非負であるようなjの最大値を考える。
A[j]→A[i]の遷移を考えると、i以下の値がA[i]-A[j]個だけP[j+1]～P[i]に入らなければならない。
j以下の未出の値を入れるか、jより大きくi以下の値を入れるかで、何個入れるかを総当たりしながら、そのパターン数を数え上げよう。

int T,N;
int A[202020];
const ll mo=998244353;
const int NUM_=400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
ll comb(ll N_, ll C_) {
	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;
	
	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;
	cin>>T;
	while(T--) {
		cin>>N;
		FOR(i,N) {
			cin>>A[i+1];
		}
		
		int ng=0;
		if(A[N]==-1) A[N]=N;
		if(A[N]!=N) ng=1;
		
		FOR(i,N+1) if(A[i]>i) ng=1;
		if(ng) {
			cout<<0<<endl;
			continue;
		}
		int pre=0;
		ll ret=1;
		for(i=1;i<=N;i++) if(A[i]!=-1) {
			if(A[pre]>A[i]) {
				ret=0;
				break;
			}
			if(A[pre]+2*(i-pre)<A[i]) {
				ret=0;
				break;
			}
			
			int fr=pre-A[pre];
			int add=A[i]-A[pre];
			ll pat=0;
			for(j=0;j<=i-pre;j++) {
				if(j>add) continue;
				int lef=add-j;
				ll a=comb(i-pre,lef)*comb(fr,lef)%mo*fact[lef]%mo;
				a*=comb(i-pre+fr-lef,j)*comb(i-pre,j)%mo*fact[j]%mo;
				(pat+=a)%=mo;
			}
			ret=ret*(pat%mo)%mo;
			
			pre=i;
		}
		cout<<ret<<endl;
	}
}

まとめ

間に合ってよかったね。