kmjp's blog

競技プログラミング参加記です

AtCoder ABC #339 (日本レジストリサービス(JPRS)プログラミングコンテスト2024) : G - Smaller Sum

なんか途中でちんたらしてしまい微妙なタイムに。
https://atcoder.jp/contests/abc339/tasks/abc339_g

問題

整数列Aが与えられる。
以下のクエリに答えよ、なお前のクエリに答えないと、次のクエリを得られない。
部分列A[L,R]に対しX以下の値の総和を求めよ。

解法

Merge-sort Treeを使う方法と平方分割の方法がある。
以下のコードは後者。
サイズDごとにAをバケットに分解し、その中で部分列の中身を昇順にソートしたものとその累積和を持つことで、各バケットにおけるX以下の値の総和をO(logD)で計算できるようにして、各クエリO(N/DlogD)程度で処理できるようにする。

int N;
const int DI=600;
ll A[402020];
vector<int> B[DI];
vector<ll> S[DI];
int Q;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	FOR(i,(202020/DI)+2) {
		FOR(j,DI) B[i].push_back(A[i*DI+j]);
		sort(ALL(B[i]));
		S[i]={0};
		FORR(b,B[i]) S[i].push_back(S[i].back()+b);
	}
	ll ret=0;
	
	cin>>Q;
	while(Q--) {
		ll TL,TR,TX;
		cin>>TL>>TR>>TX;
		int L=TL^ret;
		int R=TR^ret;
		int X=TX^ret;
		
		ret=0;
		L--;
		
		while(L%DI&&L<R) {
			if(A[L]<=X) ret+=A[L];
			L++;
		}
		while(R%DI&&L<R) {
			R--;
			if(A[R]<=X) ret+=A[R];
		}
		while(L<R) {
			x=L/DI;
			y=lower_bound(ALL(B[x]),(int)X+1)-B[x].begin();
			ret+=S[x][y];
			L+=DI;
		}
		cout<<ret<<endl;
		
	}
}

まとめ

平方分割はバケットサイズをかなり丁寧に設定しないとTLEしがちなのがちょっと苦手。