#1008634 (C++17) No.2883 K-powered Sum of Fibonacci

• 問題: No.2883 K-powered Sum of Fibonacci
• ユーザー: nonon
• 提出日時: 2024-09-08 21:30:53
• 言語: C++17 (gcc 13.3.0 + boost 1.87.0)
• 結果: AC
• 実行時間: 6 ms / 3,000 ms
• コード長: 10,136 bytes
• コンパイル時間: 2,955 ms
• コンパイル使用メモリ: 215,276 KB
• 最終ジャッジ日時: 2025-02-24 06:00:09
• ジャッジサーバーID （参考情報）: judge3 / judge2

ファイルパターン	結果
sample	AC * 3
other	AC * 40

ソースコード

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

template<long long MOD>
struct modint {
    modint() : x(0) {}
    modint(long long v) : x(v % MOD) {
        if (x < 0) x += MOD;
    }
    long long x;
    long long val() const { return x; }
    static constexpr long long mod() noexcept { return MOD; }
    friend modint operator+(modint a, modint b) { return a.x + b.x; }
    friend modint operator-(modint a, modint b) { return a.x - b.x; }
    friend modint operator*(modint a, modint b) { return a.x * b.x; }
    friend modint operator/(modint a, modint b) { return a.x * b.inv(); }
    friend modint operator+=(modint &a, modint b) { return a = a + b; }
    friend modint operator-=(modint &a, modint b) { return a = a - b; }
    friend modint operator*=(modint &a, modint b) { return a = a * b; }
    friend modint operator/=(modint &a, modint b) { return a = a / b; }
    friend bool operator==(modint a, modint b) {return a.x == b.x; }
    friend bool operator!=(modint a, modint b) {return a.x != b.x; }
    modint operator+() const { return *this; }
    modint operator-() const { return modint() - *this; }
    modint pow(long long k) const {
        modint a = x, res = 1;
        while (k > 0) {
            if (k & 1) res *= a;
            a *= a;
            k >>= 1;
        }
        return res;
    }
    modint inv() const {
        long long a = x, b = MOD;
        long long u = 1, v = 0;
        while (b > 0) {
            long long t = a / b;
            a -= t * b;
            swap(a, b);
            u -= t * v;
            swap(u, v);
        }
        return u;
    }
    modint& operator++() {
        x++;
        if (x == MOD ) x = 0;
        return *this;
    }
    modint& operator--() {
        if (x == 0) x = MOD;
        x--;
        return *this;
    }
    modint operator++(int) {
        modint res = *this;
        ++*this;
        return res;
    }
    modint operator--(int) {
        modint res = *this;
        --*this;
        return res;
    }
};

template<typename mint>
void ntt(vector<mint> &f, bool inv = false) {
    const int mod = mint::mod();
    int n = f.size();
    mint r = 3; // mod が 998244353 でないなら適切に変更すること
    if (inv) r = r.inv();
    vector<mint> g(n);
    for (int i = n / 2; i > 0; i >>= 1) {
        mint z = r.pow( (mod - 1) / (n / i) );
        mint w = 1;
        for (int j = 0; j < n; j += i << 1) {
            for (int k = 0; k < i; k++) {
                f[i + j + k] *= w;
                g[0 / 2 + j / 2 + k] = f[j + k] + f[i + j + k];
                g[n / 2 + j / 2 + k] = f[j + k] - f[i + j + k];
            }
            w *= z;
        }
        swap(f, g);
    }
    if (inv) {
        mint inv = mint(n).inv();
        for (mint &a : f) a *= inv;
    }
}

template<typename mint>
vector<mint> convolution(vector<mint> f, vector<mint> g, int d = -1) {
    int n = f.size(), m = g.size();
    if (n == 0 || m == 0) return {};
    int sz = 1;
    while (sz < n + m - 1) sz <<= 1;
    f.resize(sz); ntt(f);
    g.resize(sz); ntt(g);
    for (int i = 0; i < sz; i++) f[i] *= g[i];
    ntt(f, true);
    if (d == -1) d = n + m - 1;
    f.resize(d);
    return f;
}

#include <cassert>

template<typename mint>
struct Formal_Power_Series : vector<mint> {
    using FPS = Formal_Power_Series;
    using vector<mint>::vector;
    FPS& operator*=(mint r) {
        for (mint &x : *this) x *= r;
        return *this;
    }
    FPS& operator/=(mint r) {
        mint invr = r.inv();
        (*this) *= invr;
        return *this; 
    }
    FPS operator*(mint r) const { return FPS(*this) *= r; }
    FPS operator/(mint r) const { return FPS(*this) /= r; }
    FPS& operator+=(const FPS &f) {
        if (this->size() < f.size()) this->resize(f.size());
        for (int i = 0; i < f.size(); i++) (*this)[i] += f[i];
        return *this;
    }
    FPS& operator-=(const FPS &f) {
        if (this->size() < f.size()) this->resize(f.size());
        for (int i = 0; i < f.size(); i++) (*this)[i] -= f[i];
        return *this;
    }
    FPS& operator*=(const FPS &f) {
        vector<mint> s(this->begin(), this->end());
        vector<mint> t(f.begin(), f.end());
        s = convolution(s, t);
        this->resize(s.size());
        for (int i = 0; i < s.size(); i++) {
            (*this)[i] = s[i];
        }
        return *this;
    }
    FPS& operator/=(const FPS &f) {
        *this *= f.inv();
        return *this;
    }
    FPS& operator%=(const FPS& f) {
        *this -= this->div(f) * f;
        this->shrink();
        return *this;
    }
    FPS operator+(const FPS &f) const { return FPS(*this) += f; }
    FPS operator-(const FPS &f) const { return FPS(*this) -= f; }
    FPS operator*(const FPS &f) const { return FPS(*this) *= f; }
    FPS operator/(const FPS &f) const { return FPS(*this) /= f; }
    FPS operator%(const FPS &f) const { return FPS(*this) %= f; }

    FPS operator-() const { return FPS{} - *this; }
    FPS div(FPS f) {
        if (this->size() < f.size()) return FPS{};
        int n = this->size() - f.size() + 1;
        return (rev().pre(n) * f.rev().inv()).pre(n).rev(n);
    }
    FPS pre(int n) {
        n = min(n, int(this->size()));
        return FPS(this->begin(), this->begin() + n);
    }
    FPS rev(int n = -1) {
        FPS f(*this);
        if (n != -1) f.resize(n);
        reverse(f.begin(), f.end());
        return f;
    }
    void shrink() {
        while (this->size() && this->back() == mint(0)) {
            this->pop_back();
        }
    }
    FPS dot(FPS f) {
        int n = min(this->size(), f.size());
        FPS g(n);
        for (int i = 0; i < n; i++) {
            g[i] = (*this)[i] * f[i];
        }
        return g;
    }
    FPS diff() {
        int n = this->size();
        if (n == 0) return FPS{};
        FPS f(n - 1);
        for (int i = 1; i < n; i++) {
            f[i - 1] = (*this)[i] * i;
        }
        return f;
    }
    FPS integral() {
        int n = this->size();
        FPS f(n + 1);
        for (int i = 0; i < n; i++) {
            f[i + 1] = (*this)[i] / (i + 1);
        }
        return f;
    }
    FPS inv() {
        int n = this->size();
        FPS f(n);
        f[0] = (*this)[0].inv();
        for (int d = 1; d < n; d *= 2) {
            FPS s(2 * d), t(2 * d);
            for (int i = 0; i < min(n, 2 * d); i++) s[i] = (*this)[i];
            for (int i = 0; i < d; i++) t[i] = f[i];
            ntt(s);
            ntt(t);
            s = s.dot(t);
            ntt(s, true);
            for (int i = 0; i < d; i++) s[i] = 0;
            ntt(s);
            s = s.dot(t);
            ntt(s, true);
            for (int i = d; i < min(n, 2 * d); i++) f[i] -= s[i];
        }
        return f;
    }
    FPS exp() {
        int n = this->size();
        FPS f(1, 1);
        for (int d = 1; d < n; d *= 2) {
            FPS g = pre(2 * d);
            g[0] += 1;
            f.resize(2 * d);
            g -= f.log();
            f *= g;
            f.resize(2 * d);
        }
        f.resize(n);
        return f;
    }
    FPS log() {
        int n = this->size();
        return (diff() * inv()).pre(n - 1).integral();
    }
    FPS pow(long long k) {
        int n = this->size();
        if (k == 0) {
            FPS f(n, 0);
            f[0] = 1;
            return f;
        }
        int c = 0;
        while (c < n && (*this)[c] == 0) c++;
        if (c > (n - 1) / k) return FPS(n, 0);
        FPS f(*this);
        for (int i = 0; i + c < n; i++) f[i] = (*this)[i + c];
        f = ((f / f[0]).log() * k).exp() * f[0].pow(k);
        FPS g(n);
        for (int i = 0; i + k * c < n; i++) g[i + k * c] = f[i];
        return g;
    }
};

using mint = modint<998244353>;
using FPS = Formal_Power_Series<mint>;

vector<mint> Berlekamp_Massey(const vector<mint> &a) {
    int n = a.size();
    vector<mint> b = {1}, c = {1};
    mint y = 1;
    for (int d = 1; d <= n; d++) {
        int k = b.size(), l = c.size();
        mint x = 0;
        for (int i = 0; i < l; i++) {
            x += c[i] * a[d - l + i];
        }
        b.push_back(0);
        k++;
        if (x == 0) continue;
        mint buf = x / y;
        if (l < k) {
            vector<mint> tmp = c;
            c.insert(c.begin(), k - l, 0);
            for (int i = 0; i < k; i++) {
                c[k - i - 1] -= buf * b[k - i -1];
            }
            b = tmp;
            y = x;
        } else {
            for (int i = 0; i < k; i++) {
                c[l - i - 1] -= buf * b[k - i - 1];
            }
        }
    }
    reverse(c.begin(), c.end());
    for (mint &x : c) x = -x;
    return c;
}

mint Bostan_Mori(FPS p, FPS q, long long k) {
    mint res = 0;
    if (p.size() >= q.size()) {
        FPS r = p.div(q);
        p -= q * r;
        p.shrink();
        if (k < r.size()) res += r[k];
    }
    if (p.empty()) return res;
    p.resize( q.size() - 1 );
    auto sub = [&](const FPS &f, bool odd = 0) -> FPS {
        int n = f.size();
        if (!odd) n++;
        FPS g(n / 2);
        for (int i = odd; i < f.size(); i += 2) g[i / 2] = f[i];
        return g;
    };
    while (k) {
        FPS q2 = q;
        for (int i = 1; i < q2.size(); i += 2) q2[i] = -q2[i];
        p = sub(p * q2, k & 1);
        q = sub(q * q2);
        k /= 2;
    }
    return res + p[0];
}

mint linear_recurrence(FPS a, FPS c, long long k) {
    FPS c2(c.size() + 1);
    for (int i = 0; i < c.size(); i++) c2[i + 1] = -c[i];
    c2[0] = 1;
    return Bostan_Mori((a * c2).pre(a.size()), c2, k);
}

mint BMBM(vector<mint> x, long long k) {
    auto tmp = Berlekamp_Massey(x);
    int n = tmp.size() - 1;
    FPS a(n), c(n);
    for (int i = 0; i < n; i++) {
        a[i] = x[i];
        c[i] = tmp[i + 1];
    }
    return linear_recurrence(a, c, k);
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    long long N;
    int K;
    cin >> N >> K;
    int M = 4 * K;
    vector<mint> F(M), A(M);
    F[0] = F[1] = 1;
    A[0] = 1, A[1] = 2;
    for (int i = 2; i < M; i++) {
        F[i] = F[i - 1] + F[i - 2];
        A[i] = A[i - 1] + F[i].pow(K);
    }
    cout << BMBM(A, N - 1).val() << endl;
}