[Algorithm] C++ 백준 13725번

아이디어 요약

목표: 선형 점화식 \(A_i = \sum_{j=1}^{k} C_j A_{i-j} \pmod{104857601}\)의 \(A_N\)을 \(k \le 3\times10^4\), \(N \le 10^{18}\)에서 계산.
생성함수 관점에서 \(A(x) = \sum a_i x^i = P(x)/Q(x)\), \(Q(x)=1-\sum_{j=1}^k C_j x^j\), \(\deg P < \deg Q\).
Bostan–Mori: \([x^n] P/Q\)를 짝/홀 분리와 \(Q(-x)\)를 이용해 \(n\)을 절반씩 줄이며 구함. 한 단계당 다항식 곱셈 2회.
모듈러 104857601은 \(2^{22}\cdot 5^2+1\) 형태의 소수로 NTT 기반 다항식 곱셈이 가능. 전체 복잡도 \(O(k \log k \log N)\).
인덱스는 입력이 1-based이므로 최종 목표는 \(n = N-1\)의 계수.

C++ 풀이

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
// 42jerrykim.github.io에서 더 많은 정보를 확인 할 수 있습니다.
#include <bits/stdc++.h>
using namespace std;

static const int MOD = 104857601; // prime

inline int addmod(int a, int b) {
    int s = a + b;
    if (s >= MOD) s -= MOD;
    return s;
}
inline int submod(int a, int b) {
    int s = a - b;
    if (s < 0) s += MOD;
    return s;
}
int modpow(int a, long long e) {
    long long r = 1, x = a;
    while (e > 0) {
        if (e & 1) r = (r * x) % MOD;
        x = (x * x) % MOD;
        e >>= 1;
    }
    return (int)r;
}
int modinv(int a) { return modpow(a, MOD - 2); }

// Find primitive root of MOD (MOD prime), using prime factors of MOD-1 = 2^22 * 5^2
int primitive_root() {
    int phi = MOD - 1;
    vector<int> primes = {2, 5};
    for (int g = 2;; ++g) {
        bool ok = true;
        for (int p : primes) if (modpow(g, phi / p) == 1) { ok = false; break; }
        if (ok) return g;
    }
}

void ntt(vector<int> &a, bool invert) {
    int n = (int)a.size();
    static int g = -1;
    if (g == -1) g = primitive_root();

    for (int i = 1, j = 0; i < n; i++) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j |= bit;
        if (i < j) swap(a[i], a[j]);
    }

    for (int len = 2; len <= n; len <<= 1) {
        int wlen = modpow(g, (MOD - 1) / len);
        if (invert) wlen = modinv(wlen);
        for (int i = 0; i < n; i += len) {
            long long w = 1;
            int half = len >> 1;
            for (int j = 0; j < half; ++j) {
                int u = a[i + j];
                int v = int((w * a[i + j + half]) % MOD);
                int x = u + v; if (x >= MOD) x -= MOD;
                int y = u - v; if (y < 0) y += MOD;
                a[i + j] = x;
                a[i + j + half] = y;
                w = (w * wlen) % MOD;
            }
        }
    }
    if (invert) {
        int inv_n = modinv(n);
        for (int &x : a) x = int((1LL * x * inv_n) % MOD);
    }
}

vector<int> convolution(const vector<int> &a, const vector<int> &b) {
    if (a.empty() || b.empty()) return {};
    int need = (int)a.size() + (int)b.size() - 1;
    int n = 1; while (n < need) n <<= 1;
    vector<int> fa(n, 0), fb(n, 0);
    for (int i = 0; i < (int)a.size(); ++i) fa[i] = a[i];
    for (int i = 0; i < (int)b.size(); ++i) fb[i] = b[i];
    ntt(fa, false); ntt(fb, false);
    for (int i = 0; i < n; ++i) fa[i] = int((1LL * fa[i] * fb[i]) % MOD);
    ntt(fa, true);
    fa.resize(need);
    return fa;
}

// Bostan–Mori: compute coefficient [x^n] of P(x)/Q(x), with deg P < deg Q, Q[0] != 0
int bostan_mori(vector<int> P, vector<int> Q, long long n) {
    while (n > 0) {
        vector<int> Qm(Q.size());
        for (int i = 0; i < (int)Q.size(); ++i) Qm[i] = (i & 1) ? (MOD - Q[i]) % MOD : Q[i];
        vector<int> S = convolution(Q, Qm); // Q * Q(-x)
        vector<int> R = convolution(P, Qm); // P * Q(-x)
        vector<int> Qn((S.size() + 1) >> 1);
        for (int i = 0; i < (int)Qn.size(); ++i) Qn[i] = S[i << 1];
        vector<int> Pn;
        if ((n & 1) == 0) {
            Pn.resize((R.size() + 1) >> 1);
            for (int i = 0; i < (int)Pn.size(); ++i) Pn[i] = R[i << 1];
        } else {
            Pn.resize(R.size() >> 1);
            for (int i = 0; i < (int)Pn.size(); ++i) Pn[i] = R[(i << 1) + 1];
        }
        P.swap(Pn); Q.swap(Qn); n >>= 1;
    }
    return int(1LL * P[0] * modinv(Q[0]) % MOD);
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int k; long long N;
    if (!(cin >> k >> N)) return 0;
    vector<int> A(k), C(k);
    for (int i = 0; i < k; ++i) { long long x; cin >> x; A[i] = int(x % MOD); }
    for (int i = 0; i < k; ++i) { long long x; cin >> x; C[i] = int(x % MOD); }

    if (N <= k) { cout << A[N - 1] % MOD << '\n'; return 0; }

    // Q(x) = 1 - c1 x - c2 x^2 - ... - ck x^k
    vector<int> Q(k + 1, 0); Q[0] = 1;
    for (int i = 1; i <= k; ++i) Q[i] = (MOD - C[i - 1]) % MOD;

    // P[n] = a_n - sum_{i=1..min(n,k)} c_i * a_{n-i}
    vector<int> P(k, 0);
    for (int n = 0; n < k; ++n) {
        long long val = A[n];
        for (int i = 1; i <= n && i <= k; ++i) {
            val -= 1LL * C[i - 1] * A[n - i] % MOD;
            if (val < 0) val += MOD;
        }
        P[n] = int(val % MOD);
    }

    long long target = N - 1; // 1-based to 0-based index
    int ans = bostan_mori(P, Q, target);
    cout << ans << '\n';
    return 0;
}