AtCoder ABC #431 (トヨタシステムズプログラミングコンテスト2025) : G - One Time Swap 2

変な嵌り方してFの方に苦戦した。
https://atcoder.jp/contests/abc431/tasks/abc431_g

問題

N要素の整数列Aが与えられる。
Aのうち2要素をswapして作れる数列は重複含めてBinom(N,2)通り作れる。

このうち、クエリとして正整数Kが与えられるので、これら数列のうち辞書順でK番目のものを求めよ。

解法

swapの結果、

• Aより辞書順で小さくなるもの
• Aと同じもの
• Aより辞書順で大きくなるもの

を順に数えて行き、適宜クエリに合致する個数が登場したらクエリの解としていく。
まず最初の「Aより辞書順で小さくなるもの」のパターンでは、i＜jとなる(i,j)対に対し、A[i]とA[j]をswapする場合を考える。
iを小さい順にみて行き、その際、BITを使いA[i]＞A[j]となるjの個数を見ていこう。

「Aより辞書順で大きくなるもの」については、同様にiを大きい順に見て行き、A[i]＜A[j]となるjの個数を見て行けばよい。

int N,Q;
int A[202020];

deque<int> P[202020];

template<class V, int ME> class BIT {
public:
	V bit[1<<ME],val[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
	void set(int e,V v) { add(e,v-val[e]);}
	int lower_bound(V val) {
		V tv=0; int i,ent=0;
		for(i=ME-1;i>=0;i--) if(tv+bit[ent+(1<<i)-1]<val) tv+=bit[ent+(1<<i)-1],ent+=(1<<i);
		return ent;
	}
};
BIT<ll,20> bt;

pair<int,int> ret[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	FOR(i,N) {
		cin>>A[i];
		P[A[i]].push_front(i);
		bt.add(A[i],1);
	}
	set<pair<ll,int>> Qs;
	FOR(i,Q) {
		ll a;
		cin>>a;
		Qs.insert({a,i});
	}
	ll num=0;
	FOR(i,N) {
		while(Qs.size()) {
			ll a=Qs.begin()->first;
			if(num+bt(A[i]-1)<a) break;
			x=bt.lower_bound(a-num);
			ll b=bt(x-1);
			ret[Qs.begin()->second]={i+1,P[x][a-(num+b)-1]+1};
			Qs.erase(Qs.begin());
		}
		num+=bt(A[i]-1);
		bt.add(A[i],-1);
		P[A[i]].pop_back();
	}
	FOR(i,N) {
		P[A[i]].push_back(i);
	}
	FOR(i,N) {
		P[A[i]].pop_front();
		while(Qs.size()) {
			ll a=Qs.begin()->first;
			if(num+P[A[i]].size()<a) break;
			ret[Qs.begin()->second]={i+1,P[A[i]][0]+1};
			Qs.erase(Qs.begin());
		}
		num+=P[A[i]].size();
	}
	for(i=N-1;i>=0;i--) {
		while(Qs.size()) {
			ll a=Qs.begin()->first;
			if(num+bt(N)-bt(A[i])<a) break;
			x=bt.lower_bound(a-num+bt(A[i]));
			ll b=bt(x-1)-bt(A[i]);
			ret[Qs.begin()->second]={i+1,P[x][a-(num+b)-1]+1};
			Qs.erase(Qs.begin());
		}
		num+=bt(N)-bt(A[i]);
		bt.add(A[i],1);
		P[A[i]].push_front(i);
	}
	
	FOR(i,Q) cout<<ret[i].first<<" "<<ret[i].second<<endl;
	
}

まとめ

こちらはすんなり解けて良かった。