SRM 675 Div1 Medium LimitedMemorySeries1

長さ n (<= 5e6) の数列が与えられる。k 番目に小さい値を答えるクエリが q (<= 100) 回与えられる。
1 MB 以下、10 sec 以下でクエリの答えの総和を求める問題。

問題文にも書かれているように、メモリが足りないので数列を保持することはできない。
数列の要素の値の最大値は 1e9+7 であり、sqrt(1e9+7) 個の int 型配列を持つことができることに着目する。(実際には、(1e9+7)^(1/k) (k は小さい整数) 個の int 配列を持つことができれば十分だが、この問題では k = 2 とできる)
すると、平方分割の要領で各クエリの答えの区間を 1/sqrt(1e9+7) 倍にできる。これを 2 回繰り返すことで、答えを求めることができる。
各クエリに対して高々 2 回だけ数列を生成すれば良いので、時間計算量は O(NQ)、空間計算量はO((maxX)^1/2) となる。(下の実装では、各クエリに対して数列の生成は高々 1 回程度となっている)

ll n, x0, a, b;
vector<int> q;
#define M 31623
#define M2 (1000000007/M)
int sum[M+100], sum2[M2+100];

// Parallel binary search で解けると思ったが、log が付いたため TLE となった
/* TLE
ll f(vector<int> &c, ll lb, ll ub){
	if(c.empty()) return 0;
	if(ub-lb<=1){
		return ub*c.size();
	}
	ll md = (lb+ub+1)/2;
	ll cnt = 0;
	ll x = x0;
	rep(i, n){
		if(x<=md) cnt++;
		x = (x*a+b)%mod;
	}
	vector<int> a, b;
	int m = c.size();
	rep(i, m){
		if(cnt>=q[c[i]]) a.pb(c[i]);
		else b.pb(c[i]);
	}
	return f(a, lb, md)+f(b, md, ub);
}
*/

class LimitedMemorySeries1 {
	public:
	long long getSum(int n_, int x0_, int a_, int b_, vector <int> q_) {
		n=n_;x0=x0_;a=a_;b=b_;q=q_;
		mset(sum, 0);
		sort(all(q));
		vector<int> c(q.size());
		rep(i, q.size()) q[i]++;
		rep(i, q.size()) c[i] = i;
		ll x = x0;
		rep(i, n){
			sum[x/M2]++;
			x = (x*a+b)%mod;
		}
		rep(i, M+10) sum[i+1]+=sum[i];
		ll res = 0;
		for(int j = 0; j <= M+10; j++){
			if(q.empty()) break;
			if(q[0]>sum[j]) continue;
			mset(sum2, 0);
			ll x = x0;
			rep(i, n){
				ll t = x/M2;
				if(t<=j){
					if(t<j) sum2[0]++;
					else sum2[x%M2]++;
				}
				x = (x*a+b)%mod;
			}
			rep(i, M2+10) sum2[i+1]+=sum2[i];
			rep(i, M2+10){
				while(!q.empty()&&q[0]<=sum2[i]){
					res += j*M2+i;
					q.erase(q.begin());
				}
			}
		}
		return res;
		//return f(c, -1, mod-1);
	}
};