Description
给一个序列 \(a\) ,\(m\) 次询问,每次询问给出 \(t, k\) 。求 \(a_t + a_{t+k}+a_{t+2k}+\cdots+a_{t+pk}\) 其中 \(t+pk \leq n\) 且 \(t+(p+1)k > n\)
\(n,m \leq 300000,a_i \leq 10^9\)
Solution
对 \(k\) 即公差分块。设定一个 \(T\) 。
当 \(k > T\) 时,直接暴力算。复杂度 \(O(\frac{n}{T})\);
当 \(k \le T\) 时,对于 \(k\) 建立一个后缀和数组 \(sum\)。\(sum_i\) 表示从 \(n\) 开始往前这么跳公差 \(k\) 跳到 \(i\) 的和。它可以倒着遍历用 \(sum_i = sum_{i+k} + a_i\) 更新。复杂度 \(O(n)\)
取 \(T = \sqrt n\) 则可以预处理出所有小于 \(T\) 的 \(k\) 的 sum。复杂度 \(O(n \sqrt n)\)
但这样空间爆炸(MLE)所以开一个 sum 数组,把询问按照 \(k\) 从小到大排序。每次若 \(k>T\) 暴力;\(k \leq T\) 时重新更新 sum。由于询问中最多有 \(T\) 个不同的数 \(\leq T\)(废话) 所以更新的复杂度不会超过 \(n \sqrt n\)
所以总时间复杂度是 \(O(n \log n + n \sqrt n)\)
Code
#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int N = 300100;
int n, m, a[N];
ll sum[N], Ans[N];
struct node {
int t, k, id;
} Q[N];
inline bool cmp(node x, node y) {
return x.k == y.k ? x.t > y.t : x.k < y.k;
}
int main() {
scanf("%d", &n);
for(int i = 1; i <= n; i++)
scanf("%d", &a[i]);
int T = floor(sqrt(n));
scanf("%d", &m);
for(int i = 1; i <= m; i++) {
scanf("%d %d", &Q[i].t, &Q[i].k);
Q[i].id = i;
} sort(Q + 1, Q + m + 1, cmp); int last = n;
for(int i = 1; i <= m; i++) {
ll ans = 0;
if(Q[i].k >= T) {
for(int j = Q[i].t; j <= n; j += Q[i].k)
ans += a[j];
} else { int k = Q[i].k, t = Q[i].t;
if(Q[i].k != Q[i - 1].k) last = n;
for(int j = last; j >= t; j--) {
sum[j] = a[j];
if(j + k <= n) sum[j] += sum[j + k];
}
last = t - 1; ans = sum[t];
} Ans[Q[i].id] = ans;
}
for(int i = 1; i <= m; i++) printf("%lld\n", Ans[i]);
return 0;
}