Codeforces Round 920 (Div. 3) F. Sum of Progression 根号分治

题目链接

题目大意

        给你一个由 n 个数字组成的数组 a 。还有 q 个形式为 s,d,k 的查询。对于每个查询 q ，求元素 a [ s ]+ a [ s + d ] * 2 +⋯+ a [ s + d * ( k − 1 )] * k 的和。换句话说，对于每个查询，都需要求出数组中从 s 开始的 k 个元素的和，步长为 d ，乘以所得序列中元素的序号。

1≤ n ≤10^5,   1≤ q ≤2*10^5 ,    −10^8 ≤ a[1],...,a[n] ≤ 10^8,   1 ≤ s, d, k ≤ n ,   s+d⋅(k−1)≤n

思路

        由 s+d⋅(k−1)≤n知，d和k的数据范围相互影响，可以考虑根号分治。令m=sqrt(n), 当d>=m时，最多由n/d个数相加，即最多有sqrt(n)个数，时间复杂度为O()，可暴力求解。

        当d较小时，观察a [ s ]+ a [ s + d ] * 2 +⋯+ a [ s + d * ( k − 1 )] * k发现，这个式子相当于初始为a[ s ]，间隔为d的后缀和，而系数可通过后缀和套后缀和得到。

用sum[ i ] [ j ]表示为初始为a[ i ],间隔为j的后缀和，则sum[ i ][ j ]=sum[ i + j ][ j ] +  a[ i ];

再对sum[ i ] [ j ]求后缀和，用ssum [ i ] [ j ]表示初始为sum [ i ] [ j ]，间隔为 j的后缀和，则ssum [ i ] [ j ]=ssum [ i +j ] [ j ]+sum [ i ] [ j ]。

则所求的式子则可通过 ssum [ s ] [ d ] - ssum [ s+d*k ] [ d ] - k*sum [ s+d*k ] [ d ]求得。

code

#include<bits/stdc++.h>
#define int long long
using namespace std;

//n最大位1e5，sqrt(n)最大位317.
//一定要加一个大于317的数，因为要求后缀和，会用到i+j,要防止越界
const int N = 1e5 + 350;

int sum[N][350], ssum[N][350], a[N];

void solve() {
	int n, q;
	cin >> n >> q;
	int m = sqrt(n);

	for (int i = 1; i <= n; ++i)
		cin >> a[i];

	for (int i = n; i >= 1; --i)
		for (int j = 1; j <= m; ++j)
			sum[i][j] = sum[i + j][j] + a[i];

	for (int i = n; i >= 1; --i)
		for (int j = 1; j <= m; ++j)
			ssum[i][j] = ssum[i + j][j] + sum[i][j];

	while (q--) {
		int s, d, k;
		cin >> s >> d >> k;
		if (d > m) {
			int ans = 0;
			for (int i = 1; i <= k; ++i) {
				ans += a[s + d * (i - 1)] * i;
			}
			cout << ans << " ";
		}
		else {
			int ans = ssum[s][d] - ssum[s + d * k][d] - k * sum[s + d * k][d];
			cout << ans << " ";
		}
	}
	cout << '\n';

	//有多组检测用例，要重置
	for (int i = n; i >= 1; --i)
		for (int j = 1; j <= m; ++j)
			sum[i][j] = ssum[i][j] = 0;
}
signed main() {
	ios::sync_with_stdio(0);
	cin.tie(0), cout.tie(0);
	int t = 1;
	cin >> t;
	while (t--) solve();
	return 0;
}