yukicoder : No.1051 PQ Permutation

★3.5位でもいいかもね。
https://yukicoder.me/problems/no/1051

問題

1～NのPermutation Aと、2値P,Qが与えられる。
1～NのPermutation Bのうち、Aより辞書順で大きく、また配列中で値Pが値Qの手前に出てくるもののうち、辞書順最小のものを求めよ。

解法

配列の後ろから順に、その位置iで初めてA[i]＜B[i]となるようなケースを考える。
Aにおいてi番目以降に出てくる値をsetに入れておき高速に求められるようにしよう。

まず、B[i]＞A[i]なので、上記set中にあるA[i]より大きな最小値vを求めよう。
そのような値がない場合、適切なB[i]を取れないので次のiに進む。

• P=vの場合
  • set中にQがあるなら、B[i]=v=Pとし、B[i+1]～B[N-1]にはsetの中身を昇順に埋める。
• Q=vの場合
  • set中においてQより大きな値があるなら、vをset中でQの次に大きな値に更新し、この判定を再度損なう。
• v!=Pかつv!=Qの場合
  • PもQもsetに含まれない場合
    • A[0]～A[i-1]でPがQの手前にくるなら、B[i]=vとし、B[i+1]～B[N-1]にはsetの中身を昇順に埋める。
    • A[0]～A[i-1]でQがPの手前にくるなら、解なし、次のiに進む。
  • PもQもsetに含まれる場合
    • B[i]=vとし、B[i+1]～B[N-1]にはsetの中身をQを除いて昇順に埋める。
    • 途中Pが登場したら、以降はsetにQを追加して継続して処理する。

int N,P,Q;
int A[202020],rev[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>P>>Q;
	set<int> L,R;
	FOR(i,N) {
		cin>>A[i];
		rev[A[i]]=i;
		L.insert(i+1);
	}
	
	
	for(i=N-1;i>=0;i--) {
		L.erase(A[i]);
		R.insert(A[i]);
		auto it=R.lower_bound(A[i]+1);
		if(it==R.end()) continue;
		
		if(L.count(P)&&rev[P]<rev[Q]) {
			out:
			A[i]=*it;
			R.erase(it);
			
			for(j=i+1;j<N;j++) {
				A[j]=*R.begin();
				R.erase(R.begin());
			}
			
			
			FOR(i,N) cout<<A[i]<<" ";
			cout<<endl;
			return;
		}
		if(R.count(P)&&R.count(Q)) {
			if(*it==Q) {
				it++;
				if(it==R.end()) continue;
			}
			if(*it==P) {
				goto out;
			}
			
			A[i]=*it;
			R.erase(it);
			
			R.erase(Q);
			for(j=i+1;j<N;j++) {
				A[j]=*R.begin();
				R.erase(R.begin());
				if(A[j]==P) R.insert(Q);
			}
			FOR(i,N) cout<<A[i]<<" ";
			cout<<endl;
			return;
			
		}
		
	}
	cout<<-1<<endl;
	
}

まとめ

場合分けは大変だけど、それ以外はそこまで難しくはない。