AtCoder ARC #149 : E - Sliding Window Sort

コードはそこまで長くないんだけどね。
https://atcoder.jp/contests/arc149/tasks/arc149_e

問題

正整数N,M,Kと、N要素の整数列Aを考える。

i=0～(K-1)の順で以下を行う。

• A[i]～A[i+M-1]を昇順に並べ替える。
• ただし、添え字がNを超えるときはNで割った剰余を取ったものとする

N要素の順列Bが与えられる。
0～(N-1)の順列Aのうち、上記処理を行うとBになるものは何通りか。

解法

ソートする領域が移動するのは面倒なので、毎回A自体が左にrotateすることとする。
そうするとこの処理はAの左M要素のうち最小値がAの末尾に行き、残り(M-1)要素がAの昇順で先頭に残ると考えられる。

• 処理回数がN-M+1回未満の場合
  • ソート対象は順列の一部なので、後ろの要素は無視して考える
• 処理回数がN-M+1回以上の場合、Aの先頭に残る(M-1)要素はAの上位(M-1)要素に固定される。
  • よって残りの値がどこに入るかを考えて行けばよい。

int N,M,K;
int B[303030],A[303030];
const ll mo=998244353;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	
	cin>>N>>M>>K;
	FOR(i,N) {
		cin>>B[i];
	}
	
	if(M+K-1>=N) {
		int dif=M+K-1-N;
		FOR(i,N) A[i]=B[i];
		ZERO(B);
		FOR(i,N) {
			if(i+M-1>=N) {
				B[i]=A[(i+dif)%N];
			}
			else {
				int span=((dif)+(N-(M-1))-1-i)/(N-(M-1));
				B[i]=A[(i+1LL*span*(N-(M-1)))%N];
			}
		}
		K-=dif;
	}
	else {
		N=M+K-1;
	}
	
	vector<int> V;
	FOR(i,N) V.push_back(B[i]);
	sort(ALL(V));
	ll ret=1;
	int pre=-1;
	FOR(i,N) {
		B[i]=lower_bound(ALL(V),B[i])-V.begin();
		if(i>=N-(M-1)) {
			if(B[i]!=i) ret=0;
		}
		else if(pre<B[i]) {
			//cout<<"!"<<M<<" "<<B[i]<<endl;
			ret=ret*M%mo;
			pre=B[i];
		}
	}
	for(i=1;i<=M-1;i++) ret=ret*i%mo;
	cout<<ret<<endl;
	
}

解法

丁寧にやれば自力で解けたかもしれないが、Dの時点で残り時間がなくて厳しいな。