Sum of Progression (Codeforces 1921F)
Tags: Math, data structures

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.


Sample Inputs:
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 2

Sample Outputs:
5 1 3 
-600000000 
22 12 55 
600000000 
171 42 118 66 -108 23 2

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def f(u, v, w):
    return u << 40 ^ v << 20 ^ w

def make_graph(n, m):
    s = [0] * (n + 3)
    for i in range(0, 2 * m, 2):
        s[x[i] + 2] += 1
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * m
    for i in range(0, 2 * m, 2):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)
t = int(input())
ans = []
for _ in range(t):
    (n, q) = map(int, input().split())
    a = [0] + list(map(int, input().split()))
    x = [0] * (2 * q)
    for i in range(q):
        (s, d, k) = map(int, input().split())
        (x[2 * i], x[2 * i + 1]) = (d, f(i, s, k))
    (Q, s0) = make_graph(n, q)
    ans0 = [0] * q
    (b, c) = ([0] * (n + 1), [0] * (n + 1))
    for d in range(1, n + 1):
        if s0[d + 1] - s0[d] <= d:
            for v in range(s0[d], s0[d + 1]):
                (i, s, k) = (Q[v] >> 40, Q[v] >> 20 & 1048575, Q[v] & 1048575)
                ans1 = 0
                for j in range(k):
                    ans1 += a[s + d * j] * (j + 1)
                ans0[i] = ans1
        else:
            for i in range(1, n + 1):
                b[i] = b[max(i - d, 0)] + a[i] * (i // d + min(i % d, 1))
            for i in range(1, n + 1):
                c[i] = c[max(i - d, 0)] + a[i]
            for v in range(s0[d], s0[d + 1]):
                (i, s, k) = (Q[v] >> 40, Q[v] >> 20 & 1048575, Q[v] & 1048575)
                ans1 = b[s + d * (k - 1)] - b[max(s - d, 0)]
                x = s // d + min(s % d, 1) - 1
                ans1 -= x * (c[s + d * (k - 1)] - c[max(s - d, 0)])
                ans0[i] = ans1
    ans.append(' '.join(map(str, ans0)))
sys.stdout.write('\n'.join(ans))

Solution (Python):
import sys
input = sys.stdin.readline
K = 200

def solve():
    (n, q) = map(int, input().split())
    A = list(map(int, input().split()))
    queries = [[] for _ in range(n + 1)]
    res = [0] * q
    for i in range(q):
        (s, d, k) = map(int, input().split())
        s -= 1
        if d >= K:
            ans = 0
            for j in range(k):
                ans += (j + 1) * A[s + j * d]
            res[i] = ans
        else:
            queries[d].append((s, k, i))
    for d in range(min(K, n + 1)):
        if not queries[d]:
            continue
        P = [0] * (n + d)
        G = [0] * (n + d)
        for (i, a) in enumerate(A):
            j = i // d
            P[i + d] += a + P[i]
            G[i + d] += (j + 1) * a + G[i]
        for (s, k, i) in queries[d]:
            ns = s + (k - 1) * d
            res[i] = G[ns + d] - G[s] - s // d * (P[ns + d] - P[s])
    return ' '.join(map(str, res))
T = int(input())
out = [solve() for _ in range(T)]
print('\n'.join(map(str, out)))

Solution (Python):
import sys
input = sys.stdin.readline
K = 200

def solve():
    (n, q) = map(int, input().split())
    A = list(map(int, input().split()))
    queries = [[] for _ in range(n + 1)]
    res = [0] * q
    for i in range(q):
        (s, d, k) = map(int, input().split())
        s -= 1
        if d >= K:
            ans = 0
            for j in range(k):
                ans += (j + 1) * A[s + j * d]
            res[i] = ans
        else:
            queries[d].append((s, k, i))
    for d in range(min(K, n + 1)):
        if not queries[d]:
            continue
        P = [0] * (n + d)
        G = [0] * (n + d)
        for (i, a) in enumerate(A):
            j = i // d
            P[i + d] += a + P[i]
            G[i + d] += (j + 1) * a + G[i]
        for (s, k, i) in queries[d]:
            ns = s + (k - 1) * d
            res[i] = G[ns + d] - G[s] - s // d * (P[ns + d] - P[s])
    return ' '.join(map(str, res))
T = int(input())
out = [solve() for _ in range(T)]
print('\n'.join(map(str, out)))