【ABC179】E - Sequence Sum

atcoder.jp

解法

愚直に数列を前から順に見ていこうとすると、計算量はO(N)になり、N ≦ 10¹⁰の制約では、時間がかかりすぎてしまう。
なので、計算量を落とす方法を考える必要がある。

mod 〇〇で考えられる問題で鳩ノ巣原理が使えるというのは典型。
この問題では、数列を初項から第M+1項まで順にみていった場合、どこかで必ず同じ値を持つ項が出てくる。
この問題の漸化式では、前の項の情報しか参照していないため、「同じ値が出てきた」＝「ループが始まった」と言い換えることができる。
ループするということは、何度も同じ順番で同じ値の項が登場するということであり、無駄な計算を省くことができる。

どこからループするかさえわかれば、あとはO(1)で計算することができるため、全体計算量はO(M)になる。

計算方法

答え ＝ ループ開始までの総和 + ループ1回分の区間の総和 × ループ回数 + 最後のループの始まりから第N項まで

実装

ループの問題は実装でバグらせやすいので気をつけよう。
具体例を考えながら、実装すると考えの整理にもなってよかった。
また、そもそもループする前に第N項に到達するというパターンもあるため、注意が必要。

void solve(long long N, long long X, long long M){
    // times[i]回目にiが登場した(まだ登場していない場合、-1)
    vector<ll> times(M, -1);
    // i回目までの数列の合計値
    vector<ll> rui(M+10, 0);

    ll cur = X; // 数列のi項目
    ll ans = X;
    N--; // 初項の分を引いておく
    times[X] = 1;
    rui[1] = X;

    ll roop_times = 0;
    for (int i = 2; i <= M+1; i++) {
        if (N <= 0) {
            c(ans);
            return;
        }

        cur = cur * cur % M;

        if (times[cur] != -1) {
            // 何項でループしているか
            roop_times = i - times[cur];
            break;
        }
        times[cur] = i;
        ans += cur;
        rui[i] = ans;
        N--;
    }

    // 最後に到達した地点
    ll t = times[cur] + roop_times - 1;
    
    // 1回のループの合計値
    ll roop = rui[t] - rui[times[cur] - 1];
    // N / roop_times => ループ回数
    ans += N / roop_times * roop;

    // ループ後の余り
    N %= roop_times;
    ll amari = rui[times[cur] + N - 1] - rui[times[cur] - 1];
    ans += amari;

    cout << ans << endl;
}

典型問題かつ、バグらせやすいから、ライブラリかスニペットにしてもいいかもしれない、と思った。