yukicoder : No.877 Range ReLU Query

これはライブラリ作っておいてよかったって感じかな…。
https://yukicoder.me/problems/no/877

問題

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

• 区間[L,R]と定数Xが与えられる。sum(max(A[i]-X,0))をi∈[L,R]について求めよ。

解法

想定解と異なる方法で解いた。
以下の方法は、クエリの回答順は問わないが、数列に更新がかからない場合に使える。

まず対象の部分列をBとする。
Bが昇順にソートされているとすると、BにX以上の要素がp個あれば、解はBの末尾p個の和からX*pを引いたものである。
Bの累積和が先に計算されていれば、pをlower_boundで求めることで上記値をO(logN)で求めることができる。

これを応用する。
SegTreeの各ノードに、下位の全要素をソートした数列およびその累積和をもっておけば良い。
区間を統合する場合はマージソートの要領で結果をマージすればよい。
通常のSegTreeは大よそ初期値の設定やメモリ消費がO(N)で済むが、これはいずれもO(NlogN)かかる。

template<class V,int NV> class SegTree_1 {
public:
	vector<vector<V>> val,sum;
	static V const def=0;
	V comp(V l,V r){ return max(l,r);};
	
	SegTree_1(){val.resize(NV*2);sum.resize(NV*2);};
	V getval(int x,int y,int v,int l=0,int r=NV,int k=1) { // x<=i<y
		if(r<=x || y<=l) return 0;
		if(x<=l && r<=y) {
			int i=lower_bound(ALL(val[k]),v+1)-val[k].begin();
			return sum[k].back()-sum[k][i]-v*(val[k].size()-i);
		}
		return getval(x,y,v,l,(l+r)/2,k*2)+getval(x,y,v,(l+r)/2,r,k*2+1);
	}
	void set(int entry,V v) {
		val[entry+NV].clear();
		val[entry+NV].push_back(v);
		sum[entry+NV].clear();
		sum[entry+NV].push_back(0);
		sum[entry+NV].push_back(v);
	}
	void build() {
		for(int i=NV-1;i>=1;i--) {
			val[i].clear();
			int a=0,b=0;
			int x=i*2,y=i*2+1;
			while(a<val[x].size() || b<val[y].size()) {
				if(a==val[x].size()) {
					val[i].push_back(val[y][b++]);
				}
				else if(b==val[y].size()) {
					val[i].push_back(val[x][a++]);
				}
				else if(val[x][a]<val[y][b]) {
					val[i].push_back(val[x][a++]);
				}
				else {
					val[i].push_back(val[y][b++]);
				}
			}
			sum[i].clear();
			sum[i].push_back(0);
			FORR(v,val[i]) sum[i].push_back(sum[i].back()+v);
		}
	}
};
SegTree_1<ll,1<<18> st;

int N,Q;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	FOR(i,N) {
		cin>>x;
		st.set(i+1,x);
	}
	st.build();
	FOR(i,Q) {
		int L,R,X;
		cin>>j>>L>>R>>X;
		cout<<st.getval(L,R+1,X)<<endl;
	}
}

まとめ

元々「X以下の要素数を数え上げる」というライブラリを持っていたので、それをちょっと応用した。
こちらは次の問題で使えるね。