cplib/linear_recurrence.hpp

• View this file on GitHub
• Last update: 2025-12-24 16:17:55+09:00
• Include: #include "cplib/linear_recurrence.hpp"

Depends on

• cplib/bostan_mori.hpp

Verified with

• test/library-checker/kth_term_of_linearly_recurrent_sequence/kth_term_of_linearly_recurrent_sequence-test-00.cpp
• test/yukicoder/214/214-test-00.cpp

Code

#pragma once

#include <cassert>
#include <cstdint>
#include <ranges>
#include <vector>
#include "cplib/bostan_mori.hpp"

namespace cplib {

// d 次の線形回帰数列 a の第 n 項を求める (0-indexed)
// Bostan-Mori 法で O(d*logd*logn) または O(d^2 * logn)
// 入力: a のうち最初の d 項、線形漸化式の 1 <= i <= d 項前の係数 c_{i-1}、
//     畳み込み関数 convolution_func
template <typename T>
T linear_recurrence(const uint64_t n,
                    const std::vector<T>& a,
                    const std::vector<T>& c,
                    const auto& convolution_func) {
    if (n < a.size()) return a[n];

    const size_t d = c.size();
    assert(a.size() >= d);

    std::vector<T> q = {1};
    q.append_range(c |
                   std::views::transform([](const T& coef) { return -coef; }));
    const std::vector<T> p = convolution_func(a, q) | std::views::take(d) |
                             std::ranges::to<std::vector<T>>();
    return bostan_mori(n, p, q, convolution_func);
}

}  // namespace cplib

#line 2 "cplib/linear_recurrence.hpp"

#include <cassert>
#include <cstdint>
#include <ranges>
#include <vector>
#line 2 "cplib/bostan_mori.hpp"

#line 7 "cplib/bostan_mori.hpp"

namespace cplib {

// 有理式 p(x)/q(x) のn次の係数を求める
// deg(p) < deg(q) かつ q[0] != 0 であること
template <typename T>
T bostan_mori(uint64_t n,
              std::vector<T> p,
              std::vector<T> q,
              const auto& convolution_func) {
    while (n > 0) {
        const size_t len_q = q.size();
        std::vector<T> q_ = q;
        for (size_t i = 1; i < len_q; i += 2) {
            q_[i] *= -1;
        }
        p = convolution_func(p, q_) | std::views::drop(n % 2) |
            std::views::stride(2) | std::ranges::to<std::vector<T>>();
        q = convolution_func(q, q_) | std::views::stride(2) |
            std::ranges::to<std::vector<T>>();
        n /= 2;
    }
    assert(p.size() > 0);
    assert(q.size() > 0);
    return p[0] / q[0];
}

}  // namespace cplib
#line 8 "cplib/linear_recurrence.hpp"

namespace cplib {

// d 次の線形回帰数列 a の第 n 項を求める (0-indexed)
// Bostan-Mori 法で O(d*logd*logn) または O(d^2 * logn)
// 入力: a のうち最初の d 項、線形漸化式の 1 <= i <= d 項前の係数 c_{i-1}、
//     畳み込み関数 convolution_func
template <typename T>
T linear_recurrence(const uint64_t n,
                    const std::vector<T>& a,
                    const std::vector<T>& c,
                    const auto& convolution_func) {
    if (n < a.size()) return a[n];

    const size_t d = c.size();
    assert(a.size() >= d);

    std::vector<T> q = {1};
    q.append_range(c |
                   std::views::transform([](const T& coef) { return -coef; }));
    const std::vector<T> p = convolution_func(a, q) | std::views::take(d) |
                             std::ranges::to<std::vector<T>>();
    return bostan_mori(n, p, q, convolution_func);
}

}  // namespace cplib

Back to top page