Coeff of Rational Function (math/fps/coeff-of-rational-function.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2024-04-27 01:27:28+09:00
• Include: #include "math/fps/coeff-of-rational-function.hpp"

概要

$K$ 次多項式 $P(x), Q(x)$ に対して $\displaystyle [x^N] \frac{P(x)}{Q(x)}$ を Bostan-Mori algorithm によって計算する。

計算量

$O(K \log K \log N)$

Required by

• Kth Term of Linearly Recurrent Sequence (math/fps/kth-term-of-linearly-recurrent-sequence.hpp)

Verified with

• test/verify/yosupo-kth-term-of-linearly-recurrent-sequence.test.cpp

Code

/**
 * @brief Coeff of Rational Function
 *
 */
template <template <typename> class FPS, typename Mint>
Mint coeff_of_rational_function(FPS<Mint> P, FPS<Mint> Q, int64_t k) {
  // compute the coefficient [x^k] P/Q of rational power series
  Mint ret = 0;
  if (P.size() >= Q.size()) {
    auto R = P / Q;
    P -= R * Q;
    P.shrink();
    if (k < (int)R.size()) ret += R[k];
  }
  if (P.empty()) return ret;
  P.resize((int)Q.size() - 1);
  auto sub = [&](const FPS<Mint> &as, bool odd) {
    FPS<Mint> bs((as.size() + !odd) / 2);
    for (int i = odd; i < (int)as.size(); i += 2) bs[i >> 1] = as[i];
    return bs;
  };
  while (k > 0) {
    auto Q2(Q);
    for (int i = 1; i < (int)Q2.size(); i += 2) Q2[i] = -Q2[i];
    P = sub(P * Q2, k & 1);
    Q = sub(Q * Q2, 0);
    k >>= 1;
  }
  return ret + P[0];
}

#line 1 "math/fps/coeff-of-rational-function.hpp"
/**
 * @brief Coeff of Rational Function
 *
 */
template <template <typename> class FPS, typename Mint>
Mint coeff_of_rational_function(FPS<Mint> P, FPS<Mint> Q, int64_t k) {
  // compute the coefficient [x^k] P/Q of rational power series
  Mint ret = 0;
  if (P.size() >= Q.size()) {
    auto R = P / Q;
    P -= R * Q;
    P.shrink();
    if (k < (int)R.size()) ret += R[k];
  }
  if (P.empty()) return ret;
  P.resize((int)Q.size() - 1);
  auto sub = [&](const FPS<Mint> &as, bool odd) {
    FPS<Mint> bs((as.size() + !odd) / 2);
    for (int i = odd; i < (int)as.size(); i += 2) bs[i >> 1] = as[i];
    return bs;
  };
  while (k > 0) {
    auto Q2(Q);
    for (int i = 1; i < (int)Q2.size(); i += 2) Q2[i] = -Q2[i];
    P = sub(P * Q2, k & 1);
    Q = sub(Q * Q2, 0);
    k >>= 1;
  }
  return ret + P[0];
}

Back to top page