给定长度为n(1e5)的数组,q次询问(2e5):s,d,k,求a[s] + a[s + d] * 2 + a[s + 2 * d] * 3 + ... + a[s + (k - 1) * d]

题目

思路:根号分治,当d * d > n时,直接暴力求;否则利用前缀和求

#include 
using namespace std;
#define int long long
#define pb push_back
const int maxn = 2e5 + 5, maxm = 2e3 + 5;
int a[maxn], b[maxn];
int n;
int sqt;
vector> ask[maxn];
void solve()
{
	int q;
	cin >> n >> q;
	for(int i = 1; i <= n; i++){
		cin >> a[i];
		ask[i].clear();
	}
	vector ans(q + 5);
	for(int i = 1; i <= q; i++){
		int s, d, k;
		cin >> s >> d >> k;
		int res = 0;
		if(d * d > n){
			for(int i = 0; i < k; i++){
				int j = s + i * d;
				res += a[j] * (i + 1);
			}
			ans[i] = res;
		}
		else{
			ask[d].pb({s, k, i});
		}
	}
	// vector s0(n + 5, 0), s1(n + 5, 0);
	for(int d = 1; d <= n; d++){
		if(ask[d].empty()) continue;
		vector s0(n + 5, 0), s1(n + 5, 0);//得放在for循环里面,对于每个d,都得置零
		for(int i = 1; i <= n; i++){
			if(i > d){
				s0[i] = s0[i - d];
				s1[i] = s1[i - d];
			}
			s0[i] += a[i];//直接加起来的前缀和
			s1[i] += a[i] * ((i - 1) / d + 1);//带权值的前缀和
		}

		for(auto [s, k, id] : ask[d]){
			int res = s1[s + (k - 1) * d];
			if(s > d) res -= s1[s - d];
			int dif = (s - 1) / d + 1 - 1;
			int sum = s0[s + (k - 1) * d];
			if(s > d) sum -= s0[s - d];
			res -= sum * dif;
			ans[id] = res;
		}
	}
	for(int i = 1; i <= q; i++){
		cout << ans[i] << " \n"[i == q];
	}
}

signed main()
{
	ios::sync_with_stdio(0);
	cin.tie(0);
	int T;
	cin >> T;
	while (T--)
	{
		solve();
	}
}

你可能感兴趣的:(codeforces,算法)