多项式牛顿迭代的分块优化

草稿。。。

半在线卷积可以多叉,牛顿迭代怎么就不能多叉呢。

在本文我们利用分块得到了 sqrt(8E),inv(10E),div(10E),exp(14E),相比以前有了一定的常数提升。

本文可以看作 Faster algorithms for the square root and reciprocal of power seriesFaster exponentials of power series 的翻译。原论文还有 inv(8.66E) 和 exp(13E) 的算法,但是过于复杂了,我没有尝试。

初次写文,好多地方不规范,欢迎指教。

分块原理

对于两个长为 mrmr 多项式 f,gf, g,我们把 ff 分成 rr 块,分别记做

f[0],f[1],,f[r1]f_{[0]}, f_{[1]}, \cdots, f_{[r-1]}

对于 gg 类似,为了简化描述,我们记 X=xmX = x^m

其乘积有一些额外的项,拿 r=2r=2 举例

fg=f[0]g[0]+(f[0]g[1]+f[1]g[0])X+f[1]g[1]X2fg = f_{[0]} \ast g_{[0]} + (f_{[0]} \ast g_{[1]} + f_{[1]} \ast g_{[0]}) X + f_{[1]} \ast g_{[1]} X^2

注意

(fg)[1](f[0]g[1]+f[1]g[0])modX(fg)_{[1]} \neq \left(f_{[0]} \ast g_{[1]} + f_{[1]} \ast g_{[0]}\right) \bmod X

因为 f[0]g[0]f_{[0]} \ast g_{[0]}2m2m 项的,可以用循环卷积转回去。写全即是

(fg)[1]=(f[0]g[1]+f[1]g[0]+f[0]g[0]X)mod(X21)modX(fg)_{[1]} = \left(f_{[0]} \ast g_{[1]} + f_{[1]} \ast g_{[0]} + f_{[0]} \ast g_{[0]} \ast X\right) \bmod (X^2 - 1) \bmod X

一般的,对于 fgfg 的第 kk 块,有公式

(fg)[k]=(i=0kf[i]g[ki]+i=0k1f[i]g[k1i]X)mod(X21)modX(f g)_{[k]} = \left(\sum_{i=0}^{k} f_{[i]} \ast g_{[k-i]} + \sum_{i=0}^{k-1} f_{[i]} \ast g_{[k-1-i]} \ast X\right) \bmod (X^2 - 1) \bmod X

所以我们对应的计算过程为

F2m1(i=0kF2m(f[i])F2m(g[ki])+i=0k1F2m(f[i])F2m(g[k1i])F2m(X))\mathcal{F}_{2m}^{-1}\left(\sum_{i=0}^{k} \mathcal{F}_{2m}(f_{[i]}) \cdot \mathcal{F}_{2m}(g_{[k-i]}) + \sum_{i=0}^{k-1} \mathcal{F}_{2m}(f_{[i]}) \cdot \mathcal{F}_{2m}(g_{[k-1-i]}) \cdot \mathcal{F}_{2m}(X)\right)

F2m(X)=(1)j\mathcal{F}_{2m}(X) = (-1)^j,无需计算。即在各块的 DFT 已经计算好的情况下,我们只需要一次 IDFT 就可以算出第 kk 块。

平方根 g2=fg^2 = f

这个东西其实和半在线卷积很像,只不过 relax 最后归于一个点,而这里是一整块在迭代。

类似于普通牛顿迭代,假设 g[0],,g[k1]g_{[0]}, \cdots, g_{[k-1]} 已经给出,那么观察缺少 g[k]g_{[k]} 所造成的影响

(g[0]++g[k1]Xk1)2=f[0]++ψXk+(g[0]++g[k]Xk)2=f[0]++f[k]Xk+\begin{aligned} (g_{[0]}+ \cdots+ g_{[k-1]} X^{k-1})^2 &= f_{[0]} + \cdots + \psi X^{k} + \cdots \\ (g_{[0]}+ \cdots+ g_{[k]} X^{k})^2 &= f_{[0]} + \cdots + f_{[k]} X^{k} + \cdots \end{aligned}

即得

f[k]ψ=2g[0]g[k]modXf_{[k]} - \psi = 2 g_{[0]} \ast g_{[k]} \bmod X

具体的步骤是

  1. 给定 g[0]g_{[0]},预计算 h=(g[0])1h = (g_{[0]})^{-1},可以使用任何 6E 的算法。
  2. 计算 F2m(h)\mathcal{F}_{2m}(h),花费 2E。
    1. 假设要计算第 kk 块。
    2. 计算 F2m(g[k1])\mathcal{F}_{2m}(g_{[k-1]}),花费 2E。
    3. 计算 ψ\psi,用前述优化只需一次 IDFT,花费 2E。
    4. 计算 g[k]12h(f[k]ψ)modXg_{[k]} \gets \frac{1}{2} h \ast (f_{[k]} - \psi) \bmod X,花费 4E。

总之,根据论文估计,若我们通过分成 rr \to \infty 块计算,大概开销是 8E。

Poly Poly::sqrt() const {
    if (deg() == 1) {
        return {front().sqrt()};
    }
    const int R = 16, iv2 = qpow(2);
    int m = get_lim((deg() - 1) / R + 1);
    Poly x = cut(m).sqrt(), h = x.inv().ntt(m * 2);
    vector<Poly> ng(R);

    for (int k = 1; x.deg() < deg(); k++) {
        ng[k - 1] = x.cut(m, (k - 1) * m).ntt(m * 2);
        Poly psi(m * 2);
        for (int j = 0; j < k; j++) {
            if (j >= 1) {
                for (int i = 0; i < m; i++)
                    psi[i] -= ng[j][i] * (ng[k - j][i] + ng[k - 1 - j][i]);
                for (int i = m; i < m * 2; i++)
                    psi[i] -= ng[j][i] * (ng[k - j][i] - ng[k - 1 - j][i]);
            } else {
                for (int i = 0; i < m; i++)
                    psi[i] -= ng[j][i] * ng[k - 1 - j][i];
                for (int i = m; i < m * 2; i++)
                    psi[i] += ng[j][i] * ng[k - 1 - j][i];
            }
        }
        psi.intt(m * 2).fillZeroH(m * 2);
        for (int j = 0; j < min(m, deg() - m * k); j++)
            psi[j] += (*this)[m * k + j];
        mul(psi, h, m * 2);
        x.redeg((k + 1) * m);
        for (int i = 0; i < m; i++)
            x[m * k + i] = psi[i] * iv2;
    }
    return x.cut(deg());
}

实际测试发现,迭代写法常数对 rr 敏感,猜测是 inv 开销过大,算足 rr 次才有优势,递归在划分上比迭代更好。块数 rr 可以不是 22 的幂。

代码请看 提交

倒数 fg=1f \ast g = 1

其分块形式是

0ψ=f[0]g[k]modX0 - \psi = f_{[0]} \ast g_{[k]} \bmod X

具体的步骤是

  1. 给定 g[0]g_{[0]},而 h=(f[0])1=g[0]h = (f_{[0]})^{-1} = g_{[0]},不用计算。
    1. 假设要计算第 kk 块。
    2. 计算 F2m(g[k1])\mathcal{F}_{2m}(g_{[k-1]})F2m(f[k])\mathcal{F}_{2m}(f_{[k]}),花费 4E。
    3. 计算 ψ\psi,花费 2E。
    4. 计算 g[k]hψmodXg_{[k]} \gets -h \ast \psi \bmod X,花费 4E。
Poly Poly::inv() const { // 10E
    if (deg() == 1) {
        return {front().inv()};
    }
    const int R = 16;
    int m = get_lim((deg() - 1) / R + 1);
    Poly x = cut(m).inv();
    vector<Poly> nf(R), ng(R);
    nf[0] = cut(m).ntt(m * 2);
    for (int k = 1; x.deg() < deg(); k++) {
        nf[k] = cut(m, k * m).ntt(m * 2); // 2E
        ng[k - 1] = x.cut(m, (k - 1) * m).ntt(m * 2); // 2E
        Poly psi(m * 2);
        for (int j = 0; j < k; j++) {
            for (int i = 0; i < m; i++)
                psi[i] -= ng[j][i] * (nf[k - j][i] + nf[k - 1 - j][i]);
            for (int i = m; i < m * 2; i++)
                psi[i] -= ng[j][i] * (nf[k - j][i] - nf[k - 1 - j][i]);
        }
        psi.intt(m * 2).fillZeroH(m * 2); // 2E
        mul(psi, ng[0], m * 2); // 4E
        x.redeg((k + 1) * m);
        for (int i = 0; i < m; i++)
            x[m * k + i] = psi[i];
    }
    return x.cut(deg());
}

rr \to \infty 时,开销大概是 10E。代码请看 提交

原论文中分成 3s3s 块,前面一块方法相同,后面两块通过更优的复用做到了 26/3=8.66E26/3 = 8.66 \mathsf{E}。我还没有阅读,感兴趣的可以尝试。

商数 fg=df \ast g = d

论文里没有这个,我觉得右边加个 dd 没啥影响,试了试的确可以。

d[k]ψ=f[0]g[k]modXd_{[k]} - \psi = f_{[0]} \ast g_{[k]} \bmod X

因为 dd 不需要做 DFT,故与倒数的方法基本相同,这里不重复了。

Poly Poly::div(Poly f) const { // 10E
    if (deg() == 1) {
        return {front() * f[0].inv()};
    }
    f.redeg(deg());
    const int R = 16;
    int m = get_lim((deg() - 1) / R + 1);
    Poly x = cut(m).div(f), h = f.cut(m).inv().ntt(m * 2);

    vector<Poly> nf(R), ng(R);

    nf[0] = f.cut(m).ntt(m * 2);
    for (int k = 1; x.deg() < deg(); k++) {
        nf[k] = f.cut(m, k * m).ntt(m * 2);
        ng[k - 1] = x.cut(m, (k - 1) * m).ntt(m * 2);
        Poly psi(m * 2);
        for (int j = 0; j < k; j++) {
            for (int i = 0; i < m; i++)
                psi[i] -= ng[j][i] * (nf[k - j][i] + nf[k - 1 - j][i]);
            for (int i = m; i < m * 2; i++)
                psi[i] -= ng[j][i] * (nf[k - j][i] - nf[k - 1 - j][i]);
        }
        psi.intt(m * 2).fillZeroH(m * 2);
        for (int j = 0; j < min(m, deg() - m * k); j++)
            psi[j] += (*this)[m * k + j];
        mul(psi, h, m * 2);
        x.redeg((k + 1) * m);
        for (int i = 0; i < m; i++)
            x[m * k + i] = psi[i];
    }
    return x.cut(deg());
}

rr \to \infty 时,开销大概是 10E。代码请看 提交

指数 g=efg = e^f

指数与前者不同,其微分方程 Δg=Δfg\Delta g = \Delta f \ast g 左右都有 gg,不能直接算。

类似的,写作分块形式

(Δg)[k]ψ=g[k](Δf)[0]modX(\Delta g)_{[k]} - \psi = g_{[k]} \ast (\Delta f)_{[0]} \bmod X

为了方便起见,我们记 Δkf=XkΔ(Xkf)\Delta_{k}f = X^{-k} \Delta(X^k f),有

Δ(f[0]+f[1]X+)=(Δ0f[0])+(Δ1f[1])+\Delta (f_{[0]} + f_{[1]} X + \cdots) = (\Delta_0 f_{[0]}) + (\Delta_1 f_{[1]}) + \cdots

u=(g[0])1u = (g_{[0]})^{-1},注意到

Δu=Δexp(f[0])=uΔf[0]\Delta u = \Delta \exp (-f_{[0]}) = -u \ast \Delta f_{[0]}

因此

ϕ=uψ=u(Δg)[k]ug[k](Δf)[0]=u(Δg)[k]+g[k](Δu)=Δk(g[k]umodX)\begin{aligned} \phi = u \ast \psi &= u \ast (\Delta g)_{[k]} - u \ast g_{[k]} \ast (\Delta_f)_{[0]}\\ &= u \ast (\Delta g)_{[k]} + g_{[k]} \ast (\Delta u) \\ &= \Delta_{k}(g_{[k]} \ast u \bmod X) \end{aligned}

如此就可以更快的取出 g[k]g_{[k]} 了。

具体的计算过程:

  1. 给定 g[0]g_{[0]}
  2. 计算 u=(g[0])1u = (g_{[0]}) ^ {-1}F2m(u)\mathcal{F}_{2m}(u),花费 6E + 2E
  3. 计算 F2m((Δf)[0])\mathcal{F}_{2m}((\Delta f)_{[0]}),花费 2E。
    1. 假设要计算第 kk
    2. 计算 F2m((Δf)[k])\mathcal{F}_{2m}((\Delta f)_{[k]})F2m(g[k1])\mathcal{F}_{2m} (g_{[k-1]}),花费 2E + 2E。
    3. 计算 ψ\psi,花费 2E。
    4. 计算 ϕuψ\phi \gets u \ast \psi,花费 4E。
    5. 计算 g[k]g[0](Δk1ϕ)g_{[k]} \gets g_{[0]} \ast (\Delta_{k}^{-1} \phi),花费 4E。

rr \to \infty 时,开销大概是 14E。代码请看 提交

Poly Poly::exp() const { // 14E
    if (deg() == 1) {
        return {1};
    }
    const int S = 16;
    int m = get_lim((deg() - 1) / S + 1);
    Poly x = cut(m).exp(), u = x.inv();
    vector<Poly> nf(S), ng(S);

    Poly df = *this;
    for (int i = 0; i < df.deg(); i++)
        df[i] *= i;
    u.ntt(m * 2);
    nf[0] = df.cut(m).ntt(m * 2);

    for (int k = 1; x.deg() < deg(); k++) {
        nf[k] = df.cut(m, k * m).ntt(m * 2);
        ng[k - 1] = x.cut(m, m * (k - 1)).ntt(m * 2);
        Poly psi(m * 2);
        for (int j = 0; j < k; j++) {
            for (int i = 0; i < m; i++)
                psi[i] += ng[j][i] * (nf[k - j][i] + nf[k - 1 - j][i]);
            for (int i = m; i < m * 2; i++)
                psi[i] += ng[j][i] * (nf[k - j][i] - nf[k - 1 - j][i]);
        }
        psi.intt(m * 2).fillZeroH(m * 2);

        mul(psi, u, m * 2).fillZeroH(m * 2);
        pre_inv(m * (k + 2));
        for (int i = 0; i < m * 2; i++)
            psi[i] *= Inv[m * k + i];
        mul(psi, ng[0], m * 2).fillZeroH(m * 2);
        x.redeg((k + 1) * m);
        for (int i = 0; i < m; i++)
            x[m * k + i] = psi[i];
    }
    return x.cut(deg());
}

原论文中分成 2s2s 块,做到了 13E。我还没有阅读,感兴趣的可以尝试。

Refence

https://arxiv.org/abs/0910.1926

https://arxiv.org/abs/0911.3110

https://negiizhao.blog.uoj.ac/blog/4671

https://hly1204.github.io/library/math/formal_power_series/formal_power_series.hpp