nth_term
(src/fps/nth_term.hpp)

• View this file on GitHub
• Last update: 2024-11-09 02:16:49+09:00
• Include: #include "src/fps/nth_term.hpp"

nth_term

mint nth_term(FPS<mint> s, ll n)

斉次線形漸化式

\[a_i = \sum\limits_{j = 0}^{d - 1} c_j a_{i - 1 - j}\]

の形で表される線形回帰数列 $a$ の前 $k$ 項が与えられたときに $a_n$ を返します．

$c$ を一意に確定させるために $k \geq 2d$ が必要です．
$k < 2d$ の場合，正しく計算できていない可能性があるので注意してください．

制約

• $0 \leq n$

計算量

• $O(k^2 + k \log k \log n)$

Depends on

• berlekamp_massey (src/fps/berlekamp_massey.hpp)
• bostan_mori (src/fps/bostan_mori.hpp)
• template (src/template/template.hpp)

Verified with

• verify/yukicoder/1559.test.cpp

Code

#pragma once
#include "../template/template.hpp"
#include "./berlekamp_massey.hpp"
#include "./bostan_mori.hpp"
template <template <typename> typename FPS, typename mint>
mint nth_term(const FPS<mint>& s, const long long n) {
    assert(n >= 0);
    const FPS<mint> c = berlekamp_massey(s);
    return bostan_mori(s, c, n);
}

#line 2 "src/template/template.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using P = pair<long long, long long>;
#define rep(i, a, b) for(long long i = (a); i < (b); ++i)
#define rrep(i, a, b) for(long long i = (a); i >= (b); --i)
constexpr long long inf = 4e18;
struct SetupIO {
    SetupIO() {
        ios::sync_with_stdio(0);
        cin.tie(0);
        cout << fixed << setprecision(30);
    }
} setup_io;
#line 3 "src/fps/berlekamp_massey.hpp"
template <template <typename> typename FPS, typename mint>
FPS<mint> berlekamp_massey(const FPS<mint>& s) {
    const int n = (int)s.size();
    FPS<mint> b = {mint(-1)}, c = {mint(-1)};
    mint y = mint(1);
    for(int ed = 1; ed <= n; ++ed) {
        int l = (int)c.size(), m = (int)b.size();
        mint x = 0;
        for(int i = 0; i < l; ++i) x += c[i] * s[ed - l + i];
        b.emplace_back(0);
        ++m;
        if(x == mint(0)) continue;
        const mint freq = x / y;
        if(l < m) {
            const auto tmp = c;
            c.insert(begin(c), m - l, mint(0));
            for(int i = 0; i < m; ++i) c[m - 1 - i] -= freq * b[m - 1 - i];
            b = tmp;
            y = x;
        } else {
            for(int i = 0; i < m; ++i) c[l - 1 - i] -= freq * b[m - 1 - i];
        }
    }
    c.pop_back();
    c = c.rev();
    return c;
}
#line 3 "src/fps/bostan_mori.hpp"
template <template <typename> typename FPS, typename mint>
mint bostan_mori(const FPS<mint>& a, const FPS<mint>& c, long long k) {
    assert(k >= 0);
    if(k < (int)a.size()) return a[k];
    assert(a.size() >= c.size());
    FPS<mint> q = FPS<mint>{1} - (c << 1);
    FPS<mint> p = (a * q).pre((int)c.size());
    if(p.empty()) return 0;
    while(k > 0) {
        auto q2 = q;
        for(int i = 1; i < (int)q2.size(); i += 2) q2[i] = -q2[i];
        const auto s = p * q2;
        const auto t = q * q2;
        if(k & 1) {
            for(int i = 1; i < (int)s.size(); i += 2) p[i >> 1] = s[i];
            for(int i = 0; i < (int)t.size(); i += 2) q[i >> 1] = t[i];
        } else {
            for(int i = 0; i < (int)s.size(); i += 2) p[i >> 1] = s[i];
            for(int i = 0; i < (int)t.size(); i += 2) q[i >> 1] = t[i];
        }
        k >>= 1;
    }
    return p[0];
}
#line 5 "src/fps/nth_term.hpp"
template <template <typename> typename FPS, typename mint>
mint nth_term(const FPS<mint>& s, const long long n) {
    assert(n >= 0);
    const FPS<mint> c = berlekamp_massey(s);
    return bostan_mori(s, c, n);
}

Back to top page