yukicoder : No.1734 Decreasing Elements

うまく解くとSTLの範囲内のデータ構造で行けるね。
https://yukicoder.me/problems/no/1734

問題

整数列Aが与えられる。
以下の手順を繰り返すとき、全要素0になるまで以下の手順を何回行う必要があるか。

• A[i]＞0となる最小のiを探す。その後、AのうちA[i]以上の要素をすべてA[i]だけ引く

解法

A中で初期値が同じ値を持つ要素は、以後ずっと同じ値を持つように遷移する。
そこで、数列Bを、初期値jの要素が現在B[j]となっている、という情報を表すものとする。

初期状態として、iを順に処理して行き、B[A[i]]が正なら、BのうちB[A[i]]以上のものからB[A[i]]だけ引く、ということを行えばよい。
ただ、これだとまだO(max(A)^2)だけかかる。
そこで、Bがどのような遷移を持つか考える。
Bは一部0となっており、それ以外の箇所は、B[L...R]=[1,2,...,R-L+1]というように初期値1公差1の等差数列になっていることがわかる。
そこで、各等差数列の開始位置と長さの対の列を考える。

• B[A[i]]を求めるには、(開始位置,長さ)のタプルを昇順にした列を二分探索すればよい
• B[A[i]]以上の要素からB[A[i]]を引くには、(長さ,開始位置)のタプルのうち、長さがB[A[i]]以上の等差数列があれば2つに分割していけばよい。

ということで、(開始位置,長さ)と(長さ,開始位置)の2種類のsetを更新していけばよい。

int N;
int A[202020];

set<pair<int,int>> seg;
set<pair<int,int>> cand;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	seg.insert({-1,0});
	seg.insert({1,202020});
	seg.insert({303030,0});
	cand.insert({202020,1});
	cand.insert({0,303030});
	

	int ret=0;
	cin>>N;
	FOR(i,N) {
		cin>>x;
		
		auto it=seg.lower_bound({x+1,-1});
		it--;
		if(it->first+it->second-1<x) continue;
		ret++;
		y=x-it->first+1;
		vector<pair<int,int>> V;
		while(cand.rbegin()->first>=y) {
			V.push_back(*cand.rbegin());
			cand.erase(*cand.rbegin());
		}
		FORR2(a,b,V) {
			seg.erase({b,a});
			if(y>1) {
				seg.insert({b,y-1});
				cand.insert({y-1,b});
			}
			if(a>y) {
				seg.insert({b+y,a-y});
				cand.insert({a-y,b+y});
			}
			
		}
		
	}
	cout<<ret<<endl;
	
}

まとめ

割と短く書けて良かった。