ARC129 F問題：Takahashi's Basics in Education and Learning

問題. Takahashi's Basics in Education and Learning

初項 $A$ と公差 $B$ からなる長さ $L$ の等差数列 $s_0, s_1, \ldots, s_{L - 1}$ がある．この等差数列を一列に並べたものをひとつの整数として見たときの $M$ による剰余を答えよ．

制約： $1 \le L, A, B < 10^{18}$ ，　 $2 \le M \le 10^9$ ，　等差数列の要素はすべて $10^{18}$ 未満

解法．冪乗法

解説を参考にした．

等差数列の要素 $s_i$ を十進数表記したときの桁数を $d_i$ とすると，等差数列を一列に並べた整数 $S$ は次のように表される．
　 $S = s_{L - 1} + 10^{d_{L - 1}} (s_{L - 2} + 10^{d_{L - 2}} (s_{L - 3} + 10^{d_{L - 3}} (s_{L - 4} \cdots s_1 + 10^{d_1} (s_0))))$
　　 $= s_{L - 1} + s_{L - 2} 10^{d_{L - 1}} + s_{L - 3} 10^{d_{L - 1} + d_{L - 2}} + \cdots + s_1 10^{d_{L - 1} + d_{L - 2} + \cdots d_{2}} + s_0 10^{d_{L - 1} + d_{L - 2} + \cdots + d_1}$
ただし，この式を愚直に計算すると $O(L)$ 時間かかるので高速に計算することを考える．

桁数が $d$ となる要素をまとめて計算する．ここで，与えられた数列は等差数列なので桁数が等しい要素はもとの数列の連続した部分数列となる．また，この部分数列は等差数列と公比が $10^d$ の等比数列の積となる（等差×等比数列 - Wikipedia）．等差 x 等比数列の一般項は知られているが，ここでは連立漸化式を立てて行列の冪乗法に帰着して解く．

まず， $10^{d - 1} \le s_l, s_{l + 1}, \ldots, s_{r - 1}, s_{r} < 10^d$ となる添字 $l, r$ を求めると，
　・ $10^{d - 1} \le s_l \Leftrightarrow 10^{d - 1} \le A + B l \Leftrightarrow \frac{10^{d - 1} - A}{B} \le l$ より $l = \lceil \frac{10^{d - 1} - A}{B} \rceil$ ，
　・ $s_r < 10^{d} \Leftrightarrow A + B r \le 10^{d} - 1 \Leftrightarrow r \le \frac{10^d - 1 - A}{B}$ より $r = \lfloor \frac{10^d - 1 - A}{B} \rfloor$
となる．ただし，実際の $l$ と $r$ はそれぞれ上で計算した値と 0 の最大値と $L - 1$ の最小値をとる．

ここで， $s_l, s_{l + 1}, \ldots, s_r$ を一列に並べた整数 $S_d$ を求める． $x_0 = 0, y_0 = s_l$ を初期値として，漸化式を $x_{i + 1} = x_{i} 10^d + y_i, \, y_{i + 1} = y_i + B$ と定義すると， $x_i$ は部分数列 $s_l, s_{l + 1}, \ldots, s_r$ の 1 項目から $i$ 項目を一列に並べた整数と等しくなり， $y_i$ は部分数列の $i + 1$ 項目となる．したがって，この連立漸化式を解き $x_{r - l + 1}$ を求めると，この値は $S_d$ と等しいので答えが求まる．
連立漸化式（連立漸化式の3通りの解き方 | 高校数学の美しい物語）は漸化式を行列で表現して冪乗法で求めると $O(\log L)$ 時間で求まる．

後は，等差数列の要素が $10^{18}$ 未満なので， $d = 1, 2, \ldots, 18$ として整数 $S_d$ を求めて， $S_1, S_2, \ldots, S_{d - 1}$ の桁数を $n_d$ としたとき， $\sum_{d = 1}^{18} 10^{n_d} S_d$ の $M$ による剰余を求めればよい． $n_d$ は部分数列の左添字から分かる．

計算時間： $O(\log L)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;
using ull = unsigned long long;
using Matrix = vector<vector<ll>>;

inline ll ceil_ll(ll x, ll y) { return (x + y - 1) / y; }

Matrix mulMatrix(const Matrix &A, const Matrix &B, const ll MOD) {
    Matrix C(3, vector<ll>(3));
    for (int i = 0; i < 3; ++i) {
        for (int j = 0; j < 3; ++j) {
            for (int k = 0; k < 3; ++k) {
                (C[i][j] += (A[i][k] * B[k][j]) % MOD) %= MOD;
            }
        }
    }
    return C;
}

Matrix powMatrix(Matrix A, ll e, const ll MOD) {
    vector<vector<ll>> X(3, vector<ll>(3));
    X[0][0] = X[1][1] = X[2][2] = 1;

    for ( ; 0 < e; e >>= 1) {
        if (e & 1) X = mulMatrix(X, A, MOD);
        A = mulMatrix(A, A, MOD);
    }
    return X;
}

ll fast_pow10(ull e, const ull MOD) {
    ull a = 10, x = 1;
    for ( ; 0 < e; e >>= 1) {
        if (e & 1) (x *= a) %= MOD;
        (a *= a) %= MOD;
    }
    return x;
}

int main() {
    cin.tie(0); ios::sync_with_stdio(false);
    // cout << setprecision(8) << setiosflags(ios::fixed);

    constexpr int DIGITS = 19;
    vector<ll> pow10(DIGITS);
    pow10[0] = 1;
    for (int i = 1; i < DIGITS; ++i)
        pow10[i] = 10 * pow10[i - 1];


    ll L, A, B, M;
    cin >> L >> A >> B >> M;

    vector<ll> begin(DIGITS), num(DIGITS);
    for (ll d = 1; d < DIGITS; ++d) {
        begin[d] = max(0ll, ceil_ll(pow10[d - 1] - A, B));
        const ll end = min((pow10[d] - A - 1) / B, L - 1);

        const ll s_l = A + B * begin[d], s_r = A + B * end;
        if (pow10[d - 1] <= s_l && s_r < pow10[d] && begin[d] <= end)
            num[d] = end - begin[d] + 1;
    }

    ll ans = 0;
    for (ll d = 1; d < DIGITS; ++d) {
        if (num[d] == 0) continue;

        Matrix X(3, vector<ll>(3));
        X[0][0] = pow10[d] % M; X[0][1] = 1;
        X[1][1] = 1; X[1][2] = B % M;
        X[2][2] = 1;

        X = powMatrix(X, num[d], M);

        const ll s0 = (A + B * begin[d]) % M;
        const ll s = ((X[0][1] * s0) % M + X[0][2]) % M;

        (ans *= fast_pow10((ull)d * num[d], M)) %= M;
        (ans += s) %= M;
    }

    cout << ans << endl;

    return 0;
}

解説を見た後なのに 21WA をしたので反省．
部分数列の右添字が間違いなどの細かいミスはすぐ見つけたけど， $10^{d (r - l + 1)}$ のオーバーフローがなかなか見つからなかった．
long long だとオーバーフローで unsigned long long にしたら通った（ツライ💧）．
あとで考えると $d = 18$ のときとかに次のようなケースで落ちる．
　900000000000000000 9999999999999999 1 1000000000
実装はもう少し簡潔に書けるけど諦め．

はじめ多項式とか考えていたけど，桁数でまとめるのはなるほどで面白かった．