Polynomial Inverse
コードについての説明
体 $K$ 上の多項式 $f$ ($\mathrm{deg}(f) = n$, $f(0) \neq 0)$ を考える.
このとき $K[x] / (x^r)$ 上 $f(x)$ の乗法逆元である $f(x)^{-1}$,
つまり $f (x) \cdot f(x)^{-1} \equiv 1\ \mathrm{mod}\ x^{r}$ を満たす $f(x)^{-1}$ を求めるアルゴリズムである.
以下の実装は $\mathbb{Z} / p \mathbb{Z}$ ($p$ は素数) の場合の実装で特に素数 $p$ について
$p = k * 2^l + 1 (k > 0, \lceil \log r \rceil \le 2^l)$ を満たすこと(数論変換に乗ること)を仮定している.
こちらの参考資料が分かりやすく, それを基に実装を行った.
$f(x)^{-1}$ の存在は以下のように構成的に示すことができ, この構成方法がそのまま $\O(n + r \log r)$ のアルゴリズムとなっている.
以下の Lemma を用いて構成する.
Lemma: 多項式 $g_i$ が $f \cdot g_i \equiv 1\ \mathrm{mod}\ x^{2^i}$ を満たすなら
$f \cdot g_{i+1} \equiv 1\ \mathrm{mod}\ x^{2^{i+1}}$ を満たす多項式 $g_{i+1}$ は
$g_{i+1} \overset{\mathrm{def}}{=} \left( \left(2 g_i - f g_i^2\right) \ \mathrm{mod}\ x^{2^{i+1}} \right)$ と定めることで得られる.
今初期値 $g_0$ は $g_0 = f(0)^{-1}$ とすると $f \cdot g \equiv 1\ \mathrm{mod}\ x$ が成り立つので取れる(ちなみに可換環上でも
$f$ がモニックであれば $f(0)^{-1} = 1$ であるので取れる).
$l \overset{\mathrm{def}}{=} \lceil \log r \rceil$ と定めると上記の定理を用いて
$f \cdot g_l \equiv 1\ \mathrm{mod}\ x^{2^l}$ を満たす $g_l$ が計算でき,
$f^{-1}$ として $g_l$ を取ると $f \cdot f^{-1} \equiv 1\ \mathrm{mod}\ x^{r}$ を満たす. ゆえに構成することができた(最後に Lemma の証明を示す).
時間計算量は $i$ 回目のイテレーションで $2^{i+1}$ サイズの多項式乗算を計算するので $\sum_{i=0}^{l - 1} 2^i \log 2^i = \O( r \log r)$ (正確には $\O(n + r \log r)$) となる.
(Lemma の証明)
$f \cdot g_{i+1} \equiv 2 f g_i - f^2 g_i^2 \equiv 1 - (f g_i - 1)^2 \ \mathrm{mod}\ x^{2^{i+1}}$ と式変形でき,
仮定より $f \cdot g_i - 1 \equiv 0\ \mathrm{mod}\ x^{2^i}$ なので $f \cdot g_{i+1} \equiv 1\ \mathrm{mod}\ x^{2^{i+1}}$ が成り立つ.
(注) 便宜上長さ $r$ の配列を解となる多項式として返しているが必ずしも $r - 1$ 次多項式とは限らないことに注意する(最高次の係数が $\mathbb{Z} / p \mathbb{Z}$ 上で $0$ となる可能性があるので).
(関数)
polynomial_inverse$(a, r)$: 長さ $n$ の配列 $a$ を係数とする $n-1$ 次多項式 $f$ を与える. $f \cdot f^{-1} \equiv 1\ \mathrm{mod}\ x^{r}$ を満たす多項式 $f^{-1}$ の係数を長さ $r$ の配列として返す.
時間計算量: $\O( n + r \log r)$
コード
#define MOD 998244353
#define root 3
unsigned int add(const unsigned int x, const unsigned int y)
{
return (x + y < MOD) ? x + y : x + y - MOD;
}
unsigned int sub(const unsigned int x, const unsigned int y)
{
return (x >= y) ? (x - y) : (MOD - y + x);
}
unsigned int mul(const unsigned int x, const unsigned int y)
{
return (unsigned long long)x * y % MOD;
}
unsigned int mod_pow(unsigned int x, unsigned int n)
{
unsigned int res = 1;
while(n > 0){
if(n & 1){ res = mul(res, x); }
x = mul(x, x);
n >>= 1;
}
return res;
}
unsigned int inverse(const unsigned int x)
{
return mod_pow(x, MOD - 2);
}
void ntt(vector<int>& a, const bool rev = false)
{
unsigned int i, j, k, l, p, q, r, s;
const unsigned int size = a.size();
if(size == 1) return;
vector<int> b(size);
r = rev ? (MOD - 1 - (MOD - 1) / size) : (MOD - 1) / size;
s = mod_pow(root, r);
vector<unsigned int> kp(size / 2 + 1, 1);
for(i = 0; i < size / 2; ++i) kp[i + 1] = mul(kp[i], s);
for(i = 1, l = size / 2; i < size; i <<= 1, l >>= 1){
for(j = 0, r = 0; j < l; ++j, r += i){
for(k = 0, s = kp[i * j]; k < i; ++k){
p = a[k + r], q = a[k + r + size / 2];
b[k + 2 * r] = add(p, q);
b[k + 2 * r + i] = mul(sub(p, q), s);
}
}
swap(a, b);
}
if(rev){
s = inverse(size);
for(i = 0; i < size; i++){ a[i] = mul(a[i], s); }
}
}
vector<int> convolute(const vector<int>& a, const vector<int>& b, int asize, int bsize, int _size)
{
const int size = asize + bsize - 1;
int t = 1;
while (t < size){ t <<= 1; }
vector<int> A(t, 0), B(t, 0);
for(int i = 0; i < asize; i++){ A[i] = (a[i] < MOD) ? a[i] : (a[i] % MOD); }
for(int i = 0; i < bsize; i++){ B[i] = (b[i] < MOD) ? b[i] : (b[i] % MOD); }
ntt(A), ntt(B);
for(int i = 0; i < t; i++){ A[i] = mul(A[i], B[i]); }
ntt(A, true);
A.resize(_size);
return A;
}
vector<int> polynomial_inverse(const vector<int>& a, int r)
{
assert(a[0] != 0);
vector<int> h = {(int)inverse(a[0])};
int t = 1;
for(int i = 0; t < r; ++i){
t <<= 1;
vector<int> res = convolute(a, convolute(h, h, t / 2, t / 2, t), min((int)a.size(), t), t, t);
for(int j = 0; j < t; ++j){
res[j] = sub(0, res[j]);
if(j < t / 2) res[j] = add(res[j], mul(2, h[j]));
}
swap(h, res);
}
h.resize(r);
return h;
}
verify 用の問題
yosupo さんの library checker : Inv of Formal Power Series 提出コード