103 D Time to Raid Cowavans

問題

n項からなる数列w[i]が与えられる。
これについてq個の次のようなクエリに答えよ。

• 整数a,bが与えられる。w[a]+w[a+b]+w[a+2*b]+……の値を求める。

制約条件

n≦3*10^5
q≦3*10^5
w[i]≦10^9

方針

毎回和を愚直に計算すると最悪O(nq)かかってTLE.

bが√nより大きいときは愚直に計算して、√n以下の時は、
部分和をあらかじめ計算しておいてクエリに答えると、計算量はO(q√n)になり間に合いそう。
だがこれはMLEである。

クエリを一度全部保存しておいて、bの順に並べ替え、
同じbが変わったら部分和を全部消去してやるようにすればメモリの消費量はO(n)で済む。
これで時間・空間計算量ともに間に合うようになる。

ソースコード

int n,a[300000],q;
ll sum[300000], ans[300000];

void run()
{
	scanf("%d",&n);
	rep(i,n)scanf("%d",a+i);
	scanf("%d",&q);
	vector<pair<int,pi> > query;
	rep(i,q){
		int s,t; scanf("%d%d",&s,&t);
		query.pb(mp(t,mp(s,i)));
	}
	sort(all(query));
	rep(i,q){
		int s=query[i].second.first, t=query[i].first, id=query[i].second.second;
		if(t<550){
			if(i==0||t!=query[i-1].first){
				rep(j,t)
				for(int k=n-j-1;k>=0;k-=t)sum[k]=(k+t<n?sum[k+t]:0)+a[k];
			}
			ans[id]=sum[s-1];
		}
		else{
			ll tmp=0;
			for(int i=s-1;i<n;i+=t)tmp+=a[i];
			ans[id]=tmp;
		}
	}
	
	rep(i,q)cout<<ans[i]<<endl;
}