自分が使っている形式的冪級数ライブラリを紹介します.AtCoder Library の modint
と convolution
をベースにしています.
バグや改善できる箇所を見つけたら 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 コンテストでまだ遭遇していないこともあり,自分は持っていません).