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: