[AtCoder] ABC 179 E – Sequence Sum

2020年9月19日2020年12月13日

問題

方針

シミュレーションすると分かると思いますが、\( A_i \) の値は周期的に遷移します。剰余は \( 0 \) から \( M – 1 \) までの値を取ります。したがって、周期的になる前の \( A_i \) の値を \( v_i \)、周期的な数列を \( w_i \) とします。

\( |v| \geq N \) のとき

このとき、\( v \) の先頭から \( N \) 個の和を計算します。

\( |v| < N \) のとき

\( N \leftarrow N – |v| \) と更新し、非負整数 \(t \) を

\[ t = \left \lfloor \dfrac{N}{|w|} \right \rfloor \]

とすると、\( t \) 回 \( w \) の列が現れるので、答えは、

\[ \sum_{i=1}^{|v|}v_i + t\sum_{i = 1}^{|w|}w_i + \sum_{i=1}^{N – t|w|}w_i\]

となります。

コード

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

ll f(ll x, ll m) {
    return (x * x) % m;
}

int main() {
    ll N, X, M;
    cin >> N >> X >> M;
    set<ll> s, c;
    s.insert(X);
    ll t = X;
    vector<ll> v, w;
    v.push_back(X);
    for (int i = 0; i < N; i++) {
        t = f(t, M);
        if (s.count(t) == 1) {
            if (c.count(t) == 1) {
                break;
            }
            w.push_back(t);
            c.insert(t);
        }
        if (c.size() == 0) {
            s.insert(t);
            v.push_back(t);
        }
    }
    
    ll ans = 0;
    if (v.size() >= N) {
        for (int i = 0; i < N; i++) {
            ans += v[i];
        }
        cout << ans << "\n";
        return 0;
    }
    
    for (ll i : v) {
        ans += i;
    }
    
    ll sum = 0;
    for (ll i : w) {
        sum += i;
    }
    
    N -= v.size();
    ll m = N / c.size();
    ans += m * sum;
    ll k = N - m * c.size();
    
    for (int i = 0; i < k; i++) {
        ans += w[i];
    }
    cout << ans << "\n";
    return 0;
}

感想

周期的な問題は以前出題されたことがありました。