F. Sum of Progression
F. Sum of Progression
You are given an array $a$ of $n$ numbers. There are also $q$ queries of the form $s, d, k$.
For each query $q$, find the sum of elements $a_s + a_{s+d} \cdot 2 + \dots + a_{s + d \cdot (k - 1)} \cdot k$. In other words, for each query, it is necessary to find the sum of $k$ elements of the array with indices starting from the $s$-th, taking steps of size $d$, multiplying it by the serial number of the element in the resulting sequence.
Input
Each test consists of several testcases. The first line contains one integer $t$ ($1 \le t \le 10^4$) — the number of testcases. Next lines contain descriptions of testcases.
The first line of each testcase contains two numbers $n, q$ ($1 \le n \le 10^5, 1 \le q \le 2 \cdot 10^5$) — the number of elements in the array $a$ and the number of queries.
The second line contains $n$ integers $a_1, ... a_n$ ($-10^8 \le a_1, ..., a_n \le 10^8$) — elements of the array $a$.
The next $q$ lines each contain three integers $s$, $d$, and $k$ ($1 \le s, d, k \le n$, $s + d\cdot (k - 1) \le n$ ).
It is guaranteed that the sum of $n$ over all testcases does not exceed $10^5$, and that the sum of $q$ over all testcases does not exceed $2 \cdot 10^5 $.
Output
For each testcase, print $q$ numbers in a separate line — the desired sums, separated with space.
Example
input
5
3 3
1 1 2
1 2 2
2 2 1
1 1 2
3 1
-100000000 -100000000 -100000000
1 1 3
5 3
1 2 3 4 5
1 2 3
2 3 2
1 1 5
3 1
100000000 100000000 100000000
1 1 3
7 7
34 87 5 42 -44 66 -32
2 2 2
4 3 1
1 3 2
6 2 1
5 2 2
2 5 2
6 1 2output
5 1 3
-600000000
22 12 55
600000000
171 42 118 66 -108 23 2
解题思路
对于询问 $(s, d, k)$，对应序列的下标的公差为 $d$，首项为 $a_{(s-1) \% d + 1}$，因此 $a_s$ 是序列中的第 $l = \left\lceil \frac{s}{d} \right\rceil$ 个元素，$a_{s + d \cdot (k-1)}$ 是序列中第 $r = \left\lceil \frac{s + d \cdot (k-1)}{d} \right\rceil$ 个元素。那么答案 $a_s + 2 \cdot a_{s+d} + \dots + k \cdot a_{s + d \cdot (k - 1)}$，就可以通过前缀和得到：$$\left( l \cdot a_s + (l+1) \cdot a_{s + d} + \dots + r \cdot a_{s + d \cdot (k - 1)} \right) - (l-1) \cdot \left( a_s + a_{s + d} + \dots + a_{s + d \cdot (k - 1)} \right)$$
由于每个询问保证 $s + d\cdot (k - 1) \le n$，粗略看作是 $d \cdot k = n$，对于这种形式的表达式可以考虑根号分治。由于每次询问暴力统计答案的时间复杂度为 $O(\frac{n}{d})$，因此如果在 $d > \sqrt{n}$ 时才暴力统计那么时间复杂度就变成 $O(\sqrt{n})$，总的询问的时间复杂度就是 $O(q \sqrt{n})$。
对于 $d \leq \sqrt{n}$ 的情况，我们就可以先预处理出所有下标公差不超过 $\sqrt{n}$ 的序列的前缀和，对应的时间复杂度为 $O(n \sqrt{n})$，此时就可以通过上面提到的公式以 $O(1)$ 的复杂度算出答案。
而预处理的话我们需要计算两种前缀和，一个是带权的前缀和，用 $s_1[i][j]$ 来表示下标公差为 $i$ 的序列的前 $\left\lceil \frac{j}{d} \right\rceil$ 个带权元素的和，转移方程为 $s_1[i][j] \gets s_1[i][j - i] + \left\lceil \frac{j}{d} \right\rceil \cdot a_j$。另外一个是普通前缀和，用 $s_2[i][j]$ 来表示下标公差为 $i$ 的序列的前 $\left\lceil \frac{j}{d} \right\rceil$ 个元素的和，转移方程为 $s_2[i][j] \gets s_2[i][j - i] + a_j$。
AC 代码如下，时间复杂度为 $O(n \sqrt{n} + q \sqrt{n})$：
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 1e5 + 10, M = 310;
int a[N];
LL s1[M][N], s2[M][N];
void solve() {
int n, m;
scanf("%d %d", &n, &m);
for (int i = 1; i <= n; i++) {
scanf("%d", a + i);
}
for (int i = 1; i <= 300; i++) {
for (int j = 1; j <= n; j++) {
s1[i][j] = (j + i - 1ll) / i * a[j];
s2[i][j] = a[j];
if (j - i > 0) {
s1[i][j] += s1[i][j - i];
s2[i][j] += s2[i][j - i];
}
}
}
while (m--) {
int s, d, k;
scanf("%d %d %d", &s, &d, &k);
if (d <= 300) {
int l = max(0, s - d), r = s + d * (k - 1);
printf("%lld ", s1[d][r] - s1[d][l] - (l + d - 1) / d * (s2[d][r] - s2[d][l]));
}
else {
LL ret = 0;
for (int i = 0; i < k; i++) {
ret += (i + 1ll) * a[s + i * d];
}
printf("%lld ", ret);
}
}
printf("\n");
}
int main() {
int t;
scanf("%d", &t);
while (t--) {
solve();
}
return 0;
}再给出我一开始的做法，也是想着预处理所有的前缀和，对于每个下标公差为 $d$ 的序列有 $d$ 个不同的首项，因此总的时间复杂度就是 $O \left( \sum\limits_{i=1}^{n}{i \cdot \frac{n}{i}} \right) = O(n^2)$。但最多只有 $q$ 个询问，意味着只需计算出其中 $q$ 个前缀和即可。
因此我的做法是对于每个询问 $(s,d,k)$，记录二元组 $\left( (s-1) \% d + 1, d \right)$，如果这个二元组之前没有出现过则求出下标公差为 $d$ 首项为 $a_{(s-1) \% d + 1}$ 序列的前缀和，并把结果保存下来用编号标记，如果之后询问的二元组还是这个就可以通过 $O(1)$ 算出结果。
考虑时间复杂度和空间复杂度，最坏的情况下每次询问的二元组 $\left( (s-1) \% d + 1, d \right)$ 都不一样，并且是计算量最大的前 $q$ 个序列，即 $\frac{n}{1}, \, \frac{n}{2}, \, \frac{n}{2}, \, \frac{n}{3}, \, \ldots, \, \frac{n}{\sqrt{2q}}$。因此时间复杂度和空间复杂度都是 $O \left( \sum\limits_{i=1}^{\sqrt{2q}}{i \cdot \frac{n}{i}} \right) = O(\sqrt{q} \cdot n)$。
当 $n$ 和 $q$ 都取题目中给定的最大值时，其中所用内存大概是 $633 \times 10^5 \times 8 \times 2 \; / \; 2^{20} \approx 965 \text{ MB}$，题目给出的限制是 $1024 \text{ MB}$，然而一直 MLE。
所以优化的地方是对询问进行离线处理，把询问按照 $\left( (s-1) \% d + 1, d \right)$ 进行分组，对于公差和首项都相同的询问统一处理，结束后把空间释放，这样空间复杂度就可以降到 $O(n)$ 了。
AC 代码如下，时间复杂度为 $O(q\log{q} + \sqrt{q} \cdot n)$：
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 2e5 + 10;
int a[N];
LL s1[N], s2[N];
array<int, 3> q[N];
int p[N];
LL ans[N];
LL get(int x, int y) { // 把二元组映射成一个整数
return ((x - 1) % y + 1) * 100001ll + y;
}
void solve() {
int n, m;
scanf("%d %d", &n, &m);
for (int i = 1; i <= n; i++) {
scanf("%d", a + i);
}
for (int i = 0; i < m; i++) {
scanf("%d %d %d", &q[i][0], &q[i][1], &q[i][2]);
p[i] = i;
}
sort(p, p + m, [&](int i, int j) {
return get(q[i][0], q[i][1]) < get(q[j][0], q[j][1]);
});
for (int i = 0; i < m; i++) {
int s = q[p[i]][0], d = q[p[i]][1], k = q[p[i]][2];
for (int i = (s - 1) % d + 1, j = 1; i <= n; i += d, j++) { // 先求出前缀和
s1[j] = s1[j - 1] + 1ll * a[i] * j;
s2[j] = s2[j - 1] + a[i];
}
int j = i;
while (j < m && get(q[p[i]][0], q[p[i]][1]) == get(q[p[j]][0], q[p[j]][1])) { // 同一个序列统一处理
s = q[p[j]][0], d = q[p[j]][1], k = q[p[j]][2];
int l = (s + d - 1) / d - 1, r = (s + d * (k - 1) + d - 1) / d;
ans[p[j++]] = s1[r] - s1[l] - (s2[r] - s2[l]) * l;
}
i = j - 1;
}
for (int i = 0; i < m; i++) {
printf("%lld ", ans[i]);
}
printf("\n");
}
int main() {
int t;
scanf("%d", &t);
while (t--) {
solve();
}
return 0;
}
参考资料
Editorial for Codeforces Round 920 (Div. 3)：https://codeforces.com/blog/entry/124757
Codeforces Round 920 (Div. 3) F Sum of Progression 根号分治+前缀和：https://zhuanlan.zhihu.com/p/678014982
本文来自博客园，作者：onlyblues，转载请注明原文链接：https://www.cnblogs.com/onlyblues/p/17968882
