数式をほぼ使わないABC452-Eの解説

　ABC452 がありました。E問題 が難しかったですが、なんとか解けました。解説が数式だらけで難しいので、自分がした考察をメモしておきたいと思います。ほぼ数式なしで出来ます。

実験をする

　まず、得体の知れない $i \mod j$、これの正体を掴むために、表を書いてみます。

　縦方向に見ると、

• $0,0,0,0,0,0,0,0,0,0,\cdots$
• $1,0,1,0,1,0,1,0,1,0,\cdots$
• $1,2,0,1,2,0,1,2,0,1,\cdots$
• $1,2,3,0,1,2,3,0,1,2,\cdots$
• $1,2,3,4,0,1,2,3,4,0,\cdots$
• $1,2,3,4,5,0,1,2,3,4,\cdots$
• $1,2,3,4,5,6,0,1,2,3,\cdots$

　という風に、規則的に変化していくことがわかります。よって、$j$ を固定すればよさそうです。

　$j$ を固定し、$B_j \times (1,2, \cdots)$ を足していけば答えは出ますが、問題はこれに $A_i$ の重みが乗っかっていることです。

　例えば、$j=4$ のとき、

• $B_4 \times (1,2,3,0,1,2,3,0,1,2,\cdots)$

　となりますが、各項には、本来 $A$ の各要素が係数として掛かっています。

　見やすくするために、$A$ の各要素を $a,b,c,d,\cdots$ で表し、また、$2c=cc$ のように倍数を文字の連結で表現します。

• $B_4 \times (a,bb,ccc,0,e,ff,ggg,0,i,\cdots)$

　こうなります。

重みの計算

　よって、問題の本質は、以下のように言い換えられます。

• $j$ ごとに、$a,bb,ccc,0,e,ff,ggg,0,i,\cdots$ といった重みを計算できるか？

　図にすると以下になります。（$j=4$ のとき）

　$j=3$ のときはこうです。

　この図から、最初に

• $a,bb,ccc,dddd,eeeee,\cdots$

　なる、「すべての重み」を計算しておいて、そこから「後続の全累積和」を等間隔でいい感じに引いていけばよいことがわかります（図の着色部分）。

　ぱっと見では、$O(NM)$ の時間がかかりそうですが、$j$ が増えるとどんどん間隔が増えていくので調和級数になり、計算量は $O(N \log M)$ になります。

コード

　自分の提出を、AIによってわかりやすく整形してもらったものです。

#include <atcoder/modint>
#include <iostream>
#include <vector>
using namespace std;
using namespace atcoder;
using mint = modint998244353;

int main() {
    int n, m;
    cin >> n >> m;

    vector<int> a(n), b(m);
    for (int i = 0; i < n; i++) cin >> a[i];
    for (int j = 0; j < m; j++) cin >> b[j];

    // suffix_sum[i] = a[i] + a[i+1] + ... + a[n-1]
    vector<mint> suffix_sum(n + 1);
    for (int i = n; i > 0; i--) {
        suffix_sum[i - 1] = suffix_sum[i] + a[i - 1];
    }

    // base_weight = Σ (i+1) * a[i]  (i = 0, ..., n-1)
    // つまり 1*a[0] + 2*a[1] + ... + n*a[n-1]
    mint base_weight = 0;
    for (int i = 0; i < n; i++) {
        base_weight += mint(i + 1) * a[i];
    }

    mint ans = 0;

    for (int j = 1; j <= m; j++) {
        // base_weight から、j の倍数位置の suffix_sum を j 回分引く
        // → pos = j-1, 2j-1, 3j-1, ... の位置で suffix_sum を引く
        mint weight = base_weight;
        for (int pos = j - 1; pos < n; pos += j) {
            weight -= suffix_sum[pos] * j;
        }
        ans += weight * b[j - 1];
    }

    cout << ans.val() << endl;
}