E - You WILL Like Sigma Problem Editorial
by
AwashAmityOak
ユーザ解説
簡単のため、 \(A\) の先頭に適当な値 \(A_0\) を追加します。任意の \(j\) に対し、 \(A_0 \cdot B_j \cdot (0 \bmod j) = 0\) のため、 \(i=0\) を含めて考えても答えには影響しません。また、 \(i \gt N\) なる \(i\) について、 \(A_i = 0\) とします。
\(j\) を固定して、 \((i \bmod j)\) の値を観察してみましょう。 \((0, 1, \dots, j-1)\) が周期的に繰り返されていることがわかります。この周期ごとに、 \(A_i \cdot (i \bmod j)\) の総和を \(O(1)\) で求めることができれば、(調和級数の議論により)全体で \(O(N \log M)\) の計算量で答えが求まります。そして、実際にこれは可能です。
\(f(l, r) = \displaystyle \sum_{i=l}^{r-1} A_i \cdot (i-l)\) とします。前述した「周期ごとの \(A_i \cdot (i \bmod j)\) の総和」は \(f\) を用いて表すことができます。 \(f(l, r)\) は、
\[{ \begin{aligned} f(l, r) &= \displaystyle \sum_{i=l}^{r-1} A_i \cdot (i-l) \\ &= \displaystyle \sum_{i=l}^{r-1} (A_i \cdot i - A_i \cdot l) \\ &= \displaystyle \sum_{i=l}^{r-1} A_i \cdot i - l\sum_{i=l}^{r-1} A_i \end{aligned} }\]
と変形できます。 \(C_i = A_i \cdot i\) とおくと、 \(\displaystyle \sum_{i=l}^{r-1} A_i \cdot i\) と \(\displaystyle \sum_{i=l}^{r-1} A_i\) は、それぞれ \(C\) と \(A\) の累積和によって \(O(1)\) で求まります。そのため、 \(f(l,r)\) 自体も \(O(1)\) で求められます。
以上を踏まえて、以下のような式変形を行います。
\[{ \begin{aligned} &\displaystyle \sum_{i=1}^N \sum_{j=1}^M A_i \cdot B_j \cdot (i \bmod j) \\ = &\displaystyle \sum_{i=0}^N \sum_{j=1}^M A_i \cdot B_j \cdot (i \bmod j) \\ = &\displaystyle \sum_{j=1}^M B_j \sum_{i=0}^N A_i \cdot (i \bmod j) \\ = &\displaystyle \sum_{j=1}^M B_j \sum_{k=0}^{\lceil\frac{N+1}{j}\rceil-1} \sum_{i=0}^{j-1} A_{jk+i} \cdot i \\ = &\displaystyle \sum_{j=1}^M B_j \sum_{k=0}^{\lceil\frac{N+1}{j}\rceil-1} f(jk, jk+j) \end{aligned} }\]
この最後の式を、そのままforループとして実装すると、時間計算量 \(O(N \log M)\) でこの問題を解くことができます。
実装例
#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using mint = atcoder::modint998244353;
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int N, M;
cin >> N >> M;
// 実装上は、 B についても長さを 1 増やし、1-indexedとして扱えるようにします。
vector<mint> A(N+1), B(M+1);
for (int i = 1; i <= N; i++) {
int x;
cin >> x;
A[i] = x;
}
for (int j = 1; j <= M; j++) {
int x;
cin >> x;
B[j] = x;
}
vector<mint> C(N+1);
for (int i = 0; i <= N; i++) C[i] = A[i] * i;
// A と C の累積和
vector<mint> acc_A(N+2), acc_C(N+2);
for (int i = 0; i <= N; i++) {
acc_A[i+1] = acc_A[i] + A[i];
acc_C[i+1] = acc_C[i] + C[i];
}
auto f = [&](int l, int r) -> mint {
r = min(r, N+1); // N+1 とのminをとると、 i > N となる A_i を陽に持たなくても良いです。
return (acc_C[r] - acc_C[l]) - l * (acc_A[r] - acc_A[l]);
};
mint ans = 0;
for (int j = 1; j <= M; j++) {
mint x = 0;
for (int k = 0; k < (N+1 + j - 1) / j; k++) {
x += f(j*k, j*k + j);
}
ans += x * B[j];
}
cout << ans.val() << endl;
}
MOD = 998244353
N, M = map(int, input().split())
# 実装上は、 B についても長さを 1 増やし、1-indexedとして扱えるようにします。
A = [0] + list(map(int, input().split()))
B = [0] + list(map(int, input().split()))
C = [0]*(N+1)
for i in range(N+1):
C[i] = A[i] * i % MOD
# A と C の累積和
acc_A = [0]*(N+2)
acc_C = [0]*(N+2)
for i in range(N+1):
acc_A[i+1] = (acc_A[i] + A[i]) % MOD
acc_C[i+1] = (acc_C[i] + C[i]) % MOD
def f(l, r):
r = min(r, N+1) # N+1 とのminをとると、 i > N となる A_i を陽に持たなくても良いです。
return ((acc_C[r] - acc_C[l]) % MOD - l * (acc_A[r] - acc_A[l]) % MOD) % MOD
ans = 0
for j in range(1, M+1):
x = 0
for k in range((N+1 + j - 1) // j):
x += f(j*k, j*k + j)
x %= MOD
ans += x * B[j] % MOD;
ans %= MOD
print(ans)
posted:
last update:
