TopCoder SRM 675 Div1 Medium LimitedMemorySeries1

なんかDistributed Code Jam(ただし分散させない)やってる気分でした。
https://community.topcoder.com/stat?c=problem_statement&pm=14088

問題

n要素の整数列X[i]がある。
この整数列はX[0],a,bの3値から線形合同法の要領で計算して生成される。
またX[i]の値域は0以上M(=10^9+7)未満である。

ここでクエリ配列Q[j]が与えられる。
Xにおいて小さい方からQ[j]番目の要素をそれぞれ求め、その総和を求めよ。

なお、この問題では実行時間の制限は10秒あるが、メモリは13MBである。

解法

nの上限は500,000のため、Xを全部メモリに格納するには32bit値で20MB必要となり本問題の条件ではメモリに格納しきれない。
代わりに実行時間は非常に長いので、何度もXを計算で生成しつつ処理することを考える。

そこでバケット法＋平方分割を使おう。
まず0～(M-1)の値をM/D個の区間に区切り、X[i]のうちそれぞれの区間に入るものが何個あるかを数えていこう。
（D=√M個程度にすると良い。）

次に各クエリQ[j]について、これらM/D個の区間に登場するXの要素数を数えることで、Q[j]番目に含まれるものがどの区間にあるかを求めることができる。
例えばQ[j]番目の要素が[P,(P+D-1)]の区間に含まれることが分かったとする。
そうしたらまだXを1周計算し、うち[P,(P+D-1)]の区間に含まれるものについてはそれぞれ登場回数を数える。
そうするとQ[j]番目の値を求めることができる。

この問題では、Q[j]番目の要素が含まれるのを求めるのにO(√M)、Xを1周計算するのにO(n)かかる。
よって計算量はO(|Q|*(n+√m))。
nの方が支配的なので、Dの値は大体で十分。

以下のコードは最大ケースで8.9s程度である。

int N;
ll mo=1000000007;
int sc[101000];
const int di=30000;
int num[di+5];

class LimitedMemorySeries1 {
	public:
	long long getSum(int n, int x0, int a, int b, vector <int> Q) {
		ll ret=0,tot=0,j,nt=0;
		int qp=0,i,lo;
		sort(Q.begin(),Q.end());
		ZERO(sc);
		
		ll c=x0;
		FOR(j,n) {
			sc[c/di]++;
			c=(c*a+b)%mo;
		}
		
		for(i=lo=0;i<=mo;i+=di,lo++) {
			
			if(qp<Q.size() && Q[qp]<nt+sc[lo]) {
				ZERO(num);
				c=x0;
				FOR(j,n) {
					if(c>=i && c<i+di) num[c-i]++;
					c=(c*a+b)%mo;
				}
				
				FOR(j,di) if(num[j]) {
					while(qp<Q.size() && Q[qp]<nt+num[j]) ret += (i+j), qp++;
					nt+=num[j];
				}
			}
			else {
				nt+=sc[lo];
			}
		}
		return ret;
		
	}
}

まとめ

最初もっと重い方法を思いついたのだけど、徐々に計算量を落とせて良かった。
それにしても今回SRMっぽくない問題が多かった？