AtCoder ARC #202 : A - Merge and Increment

出だしは良かったんだけどね。
https://atcoder.jp/contests/arc202/tasks/arc202_a

問題

整数列Aが与えられる。
Aの任意の場所に任意の正整数を挿入できるとする。
以下の処理を繰り返して|A|=1とする場合、正整数の最小挿入回数は何回か。

• 同じ2値が連続する場合、それらをマージして1要素とし、1大きい値とする。

解法

5,5,5のように、同じ値が3つ続く場合、左2つをマージするか右2つをマージするか、どちらがいいかは他の箇所の値で決まる。
よってこのような場合はとりあえずランレングス圧縮した状態で考えておき、後でつじつまを合わせよう。

また、値の小さい順に考える。
今見ている値がA[i]とする。
A[i-1]とA[i+1]のいずれか小さい方と同じ値になるまで、値の挿入によりA[i]を大きくしたのち、連結しよう。
もしA[i-1]とA[i+1]が一致する場合、RLEでA[i-1]が3個あるという風に扱っておく。

int T,N;
pair<int,int> A[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		set<pair<int,int>> S;
		set<int> avail={0,N+1};
		A[0]=A[N+1]={1<<30,1};
		
		S.insert({1<<30,0});
		S.insert({1<<30,N+1});
		FOR(i,N) {
			cin>>x;
			A[i+1]={x,1};
			S.insert({x,i+1});
			avail.insert(i+1);
		}
		ll ret=0;
		while(S.size()>=3) {
			auto p=*S.begin();
			S.erase(p);
			int cur=p.second;
			avail.erase(cur);
			auto it=avail.lower_bound(cur);
			
			int nex=*it;
			int pre=*prev(it);
			
			assert(A[pre].first!=A[cur].first);
			if(A[nex].first==A[cur].first) {
				S.erase({A[nex].first,nex});
				avail.erase(nex);
				
				A[cur].second+=A[nex].second;
				S.insert({A[cur].first,cur});
				avail.insert(cur);
				continue;
			}
			if(A[cur].second>1) {
				A[cur].first++;
				ret+=A[cur].second%2;
				A[cur].second=(A[cur].second+1)/2;
				S.insert({A[cur].first,cur});
				avail.insert(cur);
				continue;
			}
			if(S.size()==2) break;
			assert(A[pre].first>A[cur].first);
			assert(A[nex].first>A[cur].first);
			int mi=min(A[nex].first,A[pre].first);
			ret+=mi-A[cur].first;
			A[cur].first=mi;
			
			S.insert({A[cur].first,cur});
			avail.insert(cur);
		}
		cout<<ret<<endl;
	}
}

まとめ

結構コード量多かったな…。