形式的冪級数(FPS)ライブラリの実装例(C++)

自分が使っている形式的冪級数ライブラリを紹介します.AtCoder Library の modintconvolution をベースにしています.

バグや改善できる箇所を見つけたら Twitter(@opt_cp)で教えていただけると嬉しいです.

別記事に形式的冪級数の練習問題を纏めています:

実装例

ライセンスは CC0 です.
参考: CC0について ― “いかなる権利も保有しない” | クリエイティブ・コモンズ・ジャパン

#define rep2(i, m, n) for (int i = (m); i < (n); ++i)
#define rep(i, n) rep2(i, 0, n)
#define drep2(i, m, n) for (int i = (m)-1; i >= (n); --i)
#define drep(i, n) drep2(i, n, 0)

template<class T>
struct FormalPowerSeries : vector<T> {
  using vector<T>::vector;
  using vector<T>::operator=;
  using F = FormalPowerSeries;

  F operator-() const {
    F res(*this);
    for (auto &e : res) e = -e;
    return res;
  }
  F &operator*=(const T &g) {
    for (auto &e : *this) e *= g;
    return *this;
  }
  F &operator/=(const T &g) {
    assert(g != T(0));
    *this *= g.inv();
    return *this;
  }
  F &operator+=(const F &g) {
    int n = (*this).size(), m = g.size();
    rep(i, min(n, m)) (*this)[i] += g[i];
    return *this;
  }
  F &operator-=(const F &g) {
    int n = (*this).size(), m = g.size();
    rep(i, min(n, m)) (*this)[i] -= g[i];
    return *this;
  }
  F &operator<<=(const int d) {
    int n = (*this).size();
    (*this).insert((*this).begin(), d, 0);
    (*this).resize(n);
    return *this;
  }
  F &operator>>=(const int d) {
    int n = (*this).size();
    (*this).erase((*this).begin(), (*this).begin() + min(n, d));
    (*this).resize(n);
    return *this;
  }
  F inv(int d = -1) const {
    int n = (*this).size();
    assert(n != 0 && (*this)[0] != 0);
    if (d == -1) d = n;
    assert(d > 0);
    F res{(*this)[0].inv()};
    while (res.size() < d) {
      int m = size(res);
      F f(begin(*this), begin(*this) + min(n, 2*m));
      F r(res);
      f.resize(2*m), internal::butterfly(f);
      r.resize(2*m), internal::butterfly(r);
      rep(i, 2*m) f[i] *= r[i];
      internal::butterfly_inv(f);
      f.erase(f.begin(), f.begin() + m);
      f.resize(2*m), internal::butterfly(f);
      rep(i, 2*m) f[i] *= r[i];
      internal::butterfly_inv(f);
      T iz = T(2*m).inv(); iz *= -iz;
      rep(i, m) f[i] *= iz;
      res.insert(res.end(), f.begin(), f.begin() + m);
    }
    return {res.begin(), res.begin() + d};
  }

  // // fast: FMT-friendly modulus only
  // F &operator*=(const F &g) {
  //   int n = (*this).size();
  //   *this = convolution(*this, g);
  //   (*this).resize(n);
  //   return *this;
  // }
  // F &operator/=(const F &g) {
  //   int n = (*this).size();
  //   *this = convolution(*this, g.inv(n));
  //   (*this).resize(n);
  //   return *this;
  // }

  // // naive
  // F &operator*=(const F &g) {
  //   int n = (*this).size(), m = g.size();
  //   drep(i, n) {
  //     (*this)[i] *= g[0];
  //     rep2(j, 1, min(i+1, m)) (*this)[i] += (*this)[i-j] * g[j];
  //   }
  //   return *this;
  // }
  // F &operator/=(const F &g) {
  //   assert(g[0] != T(0));
  //   T ig0 = g[0].inv();
  //   int n = (*this).size(), m = g.size();
  //   rep(i, n) {
  //     rep2(j, 1, min(i+1, m)) (*this)[i] -= (*this)[i-j] * g[j];
  //     (*this)[i] *= ig0;
  //   }
  //   return *this;
  // }

  // sparse
  F &operator*=(vector<pair<int, T>> g) {
    int n = (*this).size();
    auto [d, c] = g.front();
    if (d == 0) g.erase(g.begin());
    else c = 0;
    drep(i, n) {
      (*this)[i] *= c;
      for (auto &[j, b] : g) {
        if (j > i) break;
        (*this)[i] += (*this)[i-j] * b;
      }
    }
    return *this;
  }
  F &operator/=(vector<pair<int, T>> g) {
    int n = (*this).size();
    auto [d, c] = g.front();
    assert(d == 0 && c != T(0));
    T ic = c.inv();
    g.erase(g.begin());
    rep(i, n) {
      for (auto &[j, b] : g) {
        if (j > i) break;
        (*this)[i] -= (*this)[i-j] * b;
      }
      (*this)[i] *= ic;
    }
    return *this;
  }

  // multiply and divide (1 + cz^d)
  void multiply(const int d, const T c) { 
    int n = (*this).size();
    if (c == T(1)) drep(i, n-d) (*this)[i+d] += (*this)[i];
    else if (c == T(-1)) drep(i, n-d) (*this)[i+d] -= (*this)[i];
    else drep(i, n-d) (*this)[i+d] += (*this)[i] * c;
  }
  void divide(const int d, const T c) {
    int n = (*this).size();
    if (c == T(1)) rep(i, n-d) (*this)[i+d] -= (*this)[i];
    else if (c == T(-1)) rep(i, n-d) (*this)[i+d] += (*this)[i];
    else rep(i, n-d) (*this)[i+d] -= (*this)[i] * c;
  }

  T eval(const T &a) const {
    T x(1), res(0);
    for (auto e : *this) res += e * x, x *= a;
    return res;
  }

  F operator*(const T &g) const { return F(*this) *= g; }
  F operator/(const T &g) const { return F(*this) /= g; }
  F operator+(const F &g) const { return F(*this) += g; }
  F operator-(const F &g) const { return F(*this) -= g; }
  F operator<<(const int d) const { return F(*this) <<= d; }
  F operator>>(const int d) const { return F(*this) >>= d; }
  F operator*(const F &g) const { return F(*this) *= g; }
  F operator/(const F &g) const { return F(*this) /= g; }
  F operator*(vector<pair<int, T>> g) const { return F(*this) *= g; }
  F operator/(vector<pair<int, T>> g) const { return F(*this) /= g; }
};

メンバ関数リスト

$z$ を不定元とする形式的冪級数 $f, g$ に対する演算を行います.

FormalPowerSeries 構造体は,$f$ の $z^0$ から $z^N$ の係数のみ保持します.演算の結果,$z^{-1}$ 以下や $z^{N+1}$ 以上の項が発生した場合,その情報は捨てます.つまり,演算によって vector のサイズは変更されません.この仕様は他の多くの FPS ライブラリとは異なる気がするので,注意してください.

基本的な演算

  • F operator-() : $-f$ を返す
  • F &operator+=(const F &g) : $f \gets f + g$ と更新
  • F &operator-=(const F &g) : $f \gets f - g$ と更新
  • F &operator*=(const T &c) : $f \gets cf$ と更新($c$ は定数)
  • F &operator/=(const T &c) : $f \gets f/c$ と更新($c$ は定数)
  • F &operator<<=(const int d) : $f \gets f z^d$ と更新
  • F &operator>>=(const int d) : $f \gets f / z^d$ と更新

時間計算量は $O(N)$ です.

乗法の逆元

  • F inv(int d = -1) : $1/f$ を返す.$d-1$ 次の項で打ち切る.デフォルトでは $N$ 次の項で打ち切る

atcoder::convolution を使っているため,mod が NTT-friendly な場合にしか使えません.

時間計算量は $O(N \log N)$ です.

乗除算

  • F &operator*=(const F &g) : $f \gets f g$ と更新
  • F &operator/=(const F &g) : $f \gets f/g$ と更新

fast 版と naive 版の $2$ 種類あり,上のコードではともにコメントアウトしてあります.問題に応じて,必要な方のみコメントアウトを解除して使います.fast 版は mod が NTT-friendly な場合にしか使えません.

時間計算量は fast 版が $O((N+M)\log(N+M))$,naive 版が $O(NM)$ です($M$ は $g$ の次数).

スパースな乗除算

  • F &operator*=(vector<pair<int, T>> g) : $f \gets fg$ と更新
  • F &operator/=(vector<pair<int, T>> g) : $f \gets f/g$ と更新

$g$ の非ゼロ係数が少ないときに使います.$g$ は $($次数$,$ 係数$)$ の pair を並べた vector として渡します.vector は昇順ソートされていることを仮定しています.

時間計算量は $O(NK)$ です($K$ は $g$ の非ゼロ係数の数).

$(1 + c z^d)$ の乗除算

  • void multiply(const int d, const T c) : $f \gets f \times (1 + cz^d)$ と更新
  • void divide(const int d, const T c) : $f \gets f/(1 + cz^d)$ と更新

スパースな乗除算の特殊形です.この形は頻出な上,定数倍高速化ができるため別の関数にしています.

時間計算量は $O(N)$ です.

関数値評価

  • T eval(const T &a) : $f$ に $z = a$ を「代入」したときの値を返す

時間計算量は $O(N)$ です.

その他

  • F operator+(const F &g) 等: F &operator+=(const F &g) 等の結果を返す

使い方

まず上のコードを貼り,その下に

using mint = modint998244353;
using fps = FormalPowerSeries<mint>;
using sfps = vector<pair<int, mint>>;

などと書いておきます.

$f = 3 + z + 4z^2 + z^3 + 5z^4$ と $g = 2 + 7z + z^2$ に対し,積 $fg$ を計算したい場合は,

fps f = {3, 1, 4, 1, 5};
fps g = {2, 7, 1};
fps h = f * g;

とします.$fg = 6 + 23 z + 18 z^2 + 31 z^3 + 21 z^4 + 36 z^5 + 5 z^6$ なので,h の中身はこれを $4$ 次までで打ち切った {6, 23, 18, 31, 21} となります.

以下のように長めの式も書けます:

fps h = (f - (f << 2) + fps{2}) / g;

$f = 3 + z + 4z^2 + z^3 + 5z^4$ と $g = 1 - z$ に対して $f/g$ を計算したいとき,$g$ の非ゼロ係数が少ないことに着目して

fps f = {3, 1, 4, 1, 5};
sfps g = {{0, 1}, {1, -1}};
f /= g;

と書くこともできます.こうすると f の中身は {3, 4, 8, 9, 14} となります.冪級数を $1-z$ で割る操作は,係数の累積和を取ることと等価なことに注意すると,この計算結果が正しいことが分かります.

全く同じ計算を,以下のようにも書けます:

fps f = {3, 1, 4, 1, 5};
f.divide(1, -1);

こちらの方が定数倍が軽いです.

問題への適用例: TDPC: F - 準急

この問題の答えは,\[ \frac{(1-z)(z - z^K)}{1 - 2z + z^{K+1}} \] の $z^N$ の係数に等しいです.導出の詳細は 形式的冪級数(FPS)が使える AtCoder の問題リスト(随時更新) の記事を参照してください.

よって,以下のように実装できます.「スパースな乗除算」を使わないと TLE してしまうので注意してください.

// ここに FPS ライブラリを貼る

using mint = modint1000000007;
using fps = FormalPowerSeries<mint>;
using sfps = vector<pair<int, mint>>;

int main() {
  int n, k; cin >> n >> k;

  fps f = {1, -1};
  f.resize(n+1);

  f *= sfps{{1, 1}, {k, -1}};
  f /= sfps{{0, 1}, {1, -2}, {k+1, 1}};

  cout << f[n].val() << '\n';
}

補足・注意点

  • std::vector<T> を継承しています.T には AtCoder Library の modint を渡すことを想定しています.
  • 除算 $f/g$ では $g$ の定数項が非ゼロであることを要求します.
  • mod が NTT-friendly な場合にしか使えない関数があることに注意してください.
  • 乗法の逆元 $1/f$ の計算アルゴリズムは,以下の記事を参考にしました:
  • 対数 $\log f$,指数 $\exp f$,冪乗 $f^k$ 等は必要に応じて追加してください(AtCoder の rated コンテストでまだ遭遇していないこともあり,自分は持っていません).
タイトルとURLをコピーしました