yukicoder : No.1441 MErGe

最近こういう問題名多いな。
https://yukicoder.me/problems/no/1441

問題

整数列が与えられる。以下のクエリに順次答えよ。

• 区間を指定される。区間内の要素を削除し、その値の総和からなる1要素をその位置に挿入する。
• 区間を指定される。区間の総和を答える。

解法

2つの区間和を求めるBITを使い、各要素が有効かそうでないかと、有効ならその値をそれぞれ管理する。
前者のBITを二分探索することで、区間が指定されたとき、有効な要素を求めることができる。

あとは、1つ目のクエリに対しては、(L+1)番目～R番目の有効な要素を削除し、その値をL番目に加算するようにする。
2つ目のクエリに対しては、L番目～R番目の有効な要素の総和を取るようにすればよい。

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> num,sum;

int N,Q;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	num.add(0,1);
	num.add(N+1,1);
	FOR(i,N) {
		cin>>x;
		num.add(i+1,1);
		sum.add(i+1,x);
	}
	
	while(Q--) {
		int T,L,R;
		cin>>T>>L>>R;
		if(T==1) {
			int del=R-L;
			while(del--) {
				x=num.lower_bound(L+1);
				num.add(x,-1);
			}
		}
		else {
			L=num.lower_bound(L);
			R=num.lower_bound(R+1);
			cout<<sum(R)-sum(L)<<endl;
		}
	}
	
}

まとめ

このデータ構造よく使うので、setっぽいインターフェースのクラス作っておいてもいいのかもな。