導入

ニュートン法のボトルネックは FFT / IFFT です。長さ $n$ の配列を FFT or IFFT する計算量を $\mathcal{F} _ n$ とすると、ニュートン法の計算量は $k \cdot \mathcal{F} _ n + O(n)$ と書けます。面倒ですからこれ以降この $O(n)$ は省略し、$k$ を小さくすることを目指しましょう。また、$n$ は $2$ 冪であることを仮定します。一般の場合は $2$ 冪拡張を行うので、この $k$ の値が $2$ 倍になることに注意です。

ニュートン法の FFT は、精度拡張のイテレーションの中で行われます。そこで $d$ を $2$ べきとして、$\bmod x ^ d$ の解を $\bmod x ^ {2d}$ に持ち上げる操作の詳細を見てみましょう。$d$ が $1, 2, 4, \dots, n / 2$ を走査することを考えると、各イテレーションが $k \cdot \mathcal{F} _ d$ 時間であるとき、全体の計算量が $k \cdot \mathcal{F} _ n$ になることがわかります。

なお比較対象として、$n$ ($2$ 冪) 次多項式同士の積は $6 \cdot \mathcal{F} _ n$ 時間です。

FPS 逆元の素直な計算方法

$d$ を $2$ べきとして、精度 $x ^ d$ の解を $x ^ {2d}$ に持ち上げる操作の詳細を見てみましょう。形式的冪級数 $f \in k[ [ x ] ]$ と、その $\bmod x ^ d$ 逆元 $g \in k [ [ x ] ]$ が与えられているとき、$\bmod x ^ {2d}$ 逆元 $\varphi(g)$ は

$$ \varphi(g) = g \cdot ( 2 - f \cdot g ) $$

で与えられます。

次に、この計算に必要な次数のお話をしましょう。$f$ は $2d$ 次、$g$ は $d$ 次必要です。今後これを形式的冪級数ではなくて、この次数の多項式とみなしましょう。すると $\varphi(g)$ は $4d$ 次多項式ですから、FFT のバッファサイズは $4d$ あれば十分です。このサイズの FFT を $3$ 回行いますから、計算量は $12 \cdot \mathcal{F} _ n$ です。

\begin{algorithm}
\caption{FpsInvLift($f, g, d$)}
\begin{algorithmic}
\INPUT $\mathrm{deg}(f) \lt 2d, \ \mathrm{deg}(g) \lt d$
\STATE $f ← \mathcal{F} _ {4d}(f)$
\STATE $g ← \mathcal{F} _ {4d}(g)$
\FOR{$i = 0$ \TO $4d - 1$}
    \STATE $g _ i ← g _ i \cdot ( 2 -  f _ i \cdot g _ i )$
\ENDFOR
\STATE $g ← \mathcal{F} ^ {-1} _ {4d} (g)$
\RETURN $g \mathop{\mathrm{\%}} x ^ {2d}$ \COMMENT{$\mathrm{\%}$ は剰余算です}
\end{algorithmic}
\end{algorithm}

実装

まず FFT による多項式演算を使った実装一般の注意ですが、$f, g$ の後ろのほうが適切に $0$ 埋めされた状態になっていることを確認してください。そのための最も簡単な方法として、計算のたびに専用のメモリを確保するという方法がありますが、それでは格好がよろしくありませんから、工夫をしてメモリ確保をループの外に出しましょう。

変数	役割
`g`	答えを入れる場所
`f_buf`	$\mathcal{F} _ {4d}(f)$ を入れる場所

pub fn fps_inv<const P: u64>(f: &[Fp<P>], precision: usize) -> Vec<Fp<P>> {
    let fft_len_max = precision.next_power_of_two() * 2;
    let mut g = vec![fp(0); fft_len_max];
    g[0] = f[0].inv();
    let mut f_buf = vec![fp(0); fft_len_max];
    let mut fft_len = 4;
    while fft_len <= fft_len_max {
        for i in 0..f.len().min(fft_len / 2) {
            f_buf[i] = f[i];
        }
        fft(&mut f_buf[..fft_len]);
        fft(&mut g[..fft_len]);
        for i in 0..fft_len {
            g[i] = g[i] * (fp(2) - f_buf[i] * g[i]);
        }
        ifft(&mut g[..fft_len]);
        g[fft_len / 2..fft_len].fill(fp(0));
        fft_len *= 2;
    }
    g.truncate(precision);
    g
}

5/6 倍高速化

次の性質に着目することで、FFT サイズを気合いで $2d$ に縮小しましょう。その代わり FFT の回数は $5$ 回に増えますが、結局差し引き $5 / 6$ 倍になります。

$\varphi(g)$ は $[0, d[$ 次部分が $g$ と同じですから、新たに必要なのは $[d, 2d[$ 次の項だけですね
$\varphi(g)$ の代わりに $g \cdot ( 1 - f \cdot g)$ を計算しても、$[d, 2d[$ 次の項は変わりませんね
$f \cdot g = 1 \pmod { x ^ d }$ ですから、$1 - f \cdot g$ の support は $[d, 3d[$ に含まれていますね

これらを駆使すれば、次式のように次数制限をしながら計算しても良いことがわかります。

$$ \varphi(g) = g + \left( g \cdot \left(1 - f \cdot g\right) _ {[d, 2d[} \right) _ {[d, 2d[} $$

さらに式中に $2$ 回登場する多項式積 $\cdot$ は、FFT による巡回積 (つまり $\bmod \left( x ^ {2d} - 1 \right)$) で計算しても、必要な部分は巡回で汚れずキレイに保たれている事がわかります。

つまりこう計算してよいということですね。

\begin{algorithm}
\caption{FpsInvLift($f, g, d$)}
\begin{algorithmic}
\INPUT $\mathrm{deg}(f) \lt 2d, \ \mathrm{deg}(g) \lt d$
\STATE $h ← \left( 1 - f  \cdot g \right)\mathop{\mathrm{\%}} \left( x ^ {2d} - 1 \right)$ \COMMENT{ 次数 $[d, 3d[$ ですが $2d$ で巡回していい感じになりますね }
\STATE $h _ { [0, d[ } ← 0$ \COMMENT{ これで $h = \left(1 - f \cdot g\right)  _ {[d, 2d[}$ になりますね }
\STATE $h ← ( g \cdot h ) \mathop{\mathrm{\%}} \left( x ^ {2d} - 1 \right)$ \COMMENT{ 次数 $[d, 3d[$ ですが $2d$ で巡回していい感じになりますね }
\STATE $g _ {[d, 2d[} ← h _ {[d, 2d[}$
\RETURN $g$
\end{algorithmic}
\end{algorithm}

実装

これを素直に実装すると FFT が $6$ 回必要になりますけれども、factor の $g$ が共通しているため $\mathcal{F} _ {2d} (g)$ を使い回すことで $5$ 回に削減できます。

変数	役割
`g`	答えを入れる場所
`h`	途中結果を入れる場所
`g_buf`	$\mathcal{F} _ {2d}(g)$ を入れる場所

pub fn fps_inv<const P: u64>(f: &[Fp<P>], precision: usize) -> Vec<Fp<P>> {
    let fft_len_max = precision.next_power_of_two();
    let mut g = vec![fp(0); precision];
    g[0] = f[0].inv();
    let mut h = vec![fp(0); fft_len_max];
    let mut g_fft = vec![fp(0); fft_len_max];
    let mut fft_len = 2;
    while fft_len <= fft_len_max {
        if fft_len < f.len() {
            h[..fft_len].copy_from_slice(&f[..fft_len]);
        } else {
            h[..f.len()].copy_from_slice(&f[..f.len()]);
            h[f.len()..fft_len].fill(fp(0));
        }
        for i in 0..fft_len / 2 {
            g_fft[i] = g[i];
        }
        fft(&mut h[..fft_len]);
        fft(&mut g_fft[..fft_len]);
        for i in 0..fft_len {
            h[i] = fp(1) - h[i] * g_fft[i];
        }
        ifft(&mut h[..fft_len]);
        h[..fft_len / 2].fill(fp(0));
        fft(&mut h[..fft_len]);
        for i in 0..fft_len {
            h[i] = h[i] * g_fft[i];
        }
        ifft(&mut h[..fft_len]);
        g[fft_len / 2..fft_len.min(precision)]
            .copy_from_slice(&h[fft_len / 2..fft_len.min(precision)]);
        fft_len *= 2;
    }
    g
}