線形漸化式のK項目を計算する
(Src/FPS/KthTerm.hpp)

View this file on GitHub
Last update: 2024-11-13 00:58:07+09:00
Include: #include "Src/FPS/KthTerm.hpp"

概要

漸化式 $a_{n} = \sum_{i = 1}^{d} c_{i}a_{n - i}$ の $K$ 項目を計算する。

ライブラリの使い方

template <class T, class F>
T KthTerm(u64 K, std::vector<T> A, std::vector<T> C, F mult)

引数の順番注意！: K -> A -> C -> multです！

T (template)

基本的にはmintを入れることになるだろう。

コピペ禁止コンテストなどでmintを使えない場合は、中身のライブラリの演算をmodを取るように書き換える必要がある。

K

求めたい項、非負整数。

A

漸化式の始め $d$ 項 $(A_{0}, A_{1}, \dots, A_{d - 1})$ 、 $A_{d}$ 以降が混じっていても問題ないが、反対に項数が足りないとassertにひっかかる。

C

線形漸化式の係数 $(C_{1}, C_{2}, \dots, C_{d})$

mult

多項式の掛け算を行い、結果を返す関数。

[](const auto& L, const auto& R) {
    return atcoder::convolution(L, R);
}

[](const auto& L, const auto& R) {
    if (L.empty() or R.empty()) return std::vector<mint>{};
    std::vector<mint> res(L.size() + R.size() - 1);
    for (size_t i{} ; i < L.size() ; i++) for (size_t j{} ; j < R.size() ; j++) {
        res[i + j] += L[i] * R[j];
    }
    return res;
}

基本的には以上のどちらかで事足りるハズ。

計算量解析

multを $\Theta(\log K)$ 回呼び出すのがボトルネック

multで前者の例を使うと $\Theta (d\log d\log K)$
multで後者の例を使うと $\Theta (d^2 \log K)$

補足

内部でBostan-Moriアルゴリズムが走っている。

線形漸化式の $K$ 番目を求めるアルゴリズムとして、

\[\begin{pmatrix} 0, 1, 0, \cdots, 0 \\ 0, 0, 1, \cdots, 0 \\ \vdots \\ 0, 0, 0, \cdots, 1 \\ C_{d}, C_{d - 1}, C_{d - 2}, \cdots, C{1} \\ \end{pmatrix}\]

の $K$ 乗を計算する方法(行列累乗)が有名だが、これの時間計算量は $\Theta (d^3 \log K)$ と見積もれる。

このライブラリのmultに愚直を使ったやつよりも計算量が劣っていることに注意(貼れるときはちゃんと貼ろう)

勉強不足故、定数倍の良い実装の方針をまだ採用していない。定数倍バトルになったらおとなしく別ライブラリを探そう…

Depends on

標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Verified with

Code

#pragma once

#include <cassert>
#include <vector>

#include "../Template/TypeAlias.hpp"

namespace zawa {

namespace internal {

// [x^K] P(x) / Q(x)を計算するアルゴリズム
template <class T, class F>
T BostanMori(u64 K, std::vector<T> P, std::vector<T> Q, F mult) {
    assert(P.size());
    assert(Q.size() and Q[0] != T(0));

    // p(-x)を計算
    auto minus_x{[](const std::vector<T>& p) -> std::vector<T> {
        std::vector<T> res(p.size()); 
        for (usize i{} ; i < p.size() ; i++) res[i] = (i % 2 ? T{-1} * p[i] : p[i]);
        return res;
    }};
    // 奇数次数の係数のみを取り出す
    auto odd{[](const std::vector<T>& p) -> std::vector<T> {
        std::vector<T> res;
        res.reserve(p.size() >> 1);
        for (usize i{1} ; i < p.size() ; i += 2u) res.push_back(p[i]);
        return res;
    }};
    // 偶数次数の係数のみを取り出す
    auto even{[](const std::vector<T>& p) -> std::vector<T> {
        std::vector<T> res;
        res.reserve((p.size() & 1) + (p.size() >> 1));
        for (usize i{} ; i < p.size() ; i += 2u) res.push_back(p[i]);
        return res;
    }};

    while (K) {
        auto Qm{minus_x(Q)};
        auto U{mult(P, Qm)};
        P = (K & 1 ? odd(U) : even(U));
        Q = even(mult(Q, Qm));
        K >>= 1;
    }
    return (Q[0] == T{1} ? P[0] : P[0] / Q[0]);
}

} // namespace internal

template <class T, class F>
T KthTerm(u64 K, std::vector<T> A, std::vector<T> C, F mult) {
    assert(A.size() + 1u >= C.size());
    if (K < A.size()) return A[K];
    std::vector<T> tmp(C.size() + 1, T{1});
    for (usize i{} ; i < C.size() ; i++) {
        tmp[i + 1] = -C[i];
    }
    C = std::move(tmp);
    A.resize(C.size() - 1);
    A = mult(A, C);
    A.resize(C.size() - 1);
    return internal::BostanMori(K, A, C, mult);
}

} // namespace zawa

#line 2 "Src/FPS/KthTerm.hpp"

#include <cassert>
#include <vector>

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 7 "Src/FPS/KthTerm.hpp"

namespace zawa {

namespace internal {

// [x^K] P(x) / Q(x)を計算するアルゴリズム
template <class T, class F>
T BostanMori(u64 K, std::vector<T> P, std::vector<T> Q, F mult) {
    assert(P.size());
    assert(Q.size() and Q[0] != T(0));

    // p(-x)を計算
    auto minus_x{[](const std::vector<T>& p) -> std::vector<T> {
        std::vector<T> res(p.size()); 
        for (usize i{} ; i < p.size() ; i++) res[i] = (i % 2 ? T{-1} * p[i] : p[i]);
        return res;
    }};
    // 奇数次数の係数のみを取り出す
    auto odd{[](const std::vector<T>& p) -> std::vector<T> {
        std::vector<T> res;
        res.reserve(p.size() >> 1);
        for (usize i{1} ; i < p.size() ; i += 2u) res.push_back(p[i]);
        return res;
    }};
    // 偶数次数の係数のみを取り出す
    auto even{[](const std::vector<T>& p) -> std::vector<T> {
        std::vector<T> res;
        res.reserve((p.size() & 1) + (p.size() >> 1));
        for (usize i{} ; i < p.size() ; i += 2u) res.push_back(p[i]);
        return res;
    }};

    while (K) {
        auto Qm{minus_x(Q)};
        auto U{mult(P, Qm)};
        P = (K & 1 ? odd(U) : even(U));
        Q = even(mult(Q, Qm));
        K >>= 1;
    }
    return (Q[0] == T{1} ? P[0] : P[0] / Q[0]);
}

} // namespace internal

template <class T, class F>
T KthTerm(u64 K, std::vector<T> A, std::vector<T> C, F mult) {
    assert(A.size() + 1u >= C.size());
    if (K < A.size()) return A[K];
    std::vector<T> tmp(C.size() + 1, T{1});
    for (usize i{} ; i < C.size() ; i++) {
        tmp[i + 1] = -C[i];
    }
    C = std::move(tmp);
    A.resize(C.size() - 1);
    A = mult(A, C);
    A.resize(C.size() - 1);
    return internal::BostanMori(K, A, C, mult);
}

} // namespace zawa

線形漸化式のK項目を計算する (Src/FPS/KthTerm.hpp)

概要