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

2020/11/08

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

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

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

形式的冪級数（FPS）が使える AtCoder の問題リスト（随時更新）

考察に使える問題と，実装に使える問題の両方を集めています．一部の問題には簡単な解説もつけています．

opt-cp.com
2020/10/13

実装例

ライセンスは 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+cz^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をコピーしました