$\newcommand{\O}{\mathrm{O}}$ My Algorithm : kopricky アルゴリズムライブラリ

kopricky アルゴリズムライブラリ

Polynomial Exponential

コードについての説明

体 $K$ 上の多項式 $f$ ($\mathrm{deg}(f) = n$, $f(0) = 0)$ に対応する形式的べき級数 $F$ が与えられたときに 形式的べき級数 $\exp (F)$ の $x^i (0 \le i < r)$ の係数を返す.
以下の実装は $K = \mathbb{Z} / p \mathbb{Z}$ ($p$ は素数) の場合の実装で特に素数 $p$ について $p = k * 2^l + 1 (k > 0, n + 1 \le 2^l)$ を満たすこと(数論変換に乗ること)を仮定している.
形式的べき級数の $\exp$ の求め方がよく分からなかったので あまり良くない気もするが verify 問題の latte0199 さんのコードを参照させていただいた.
その上で逆手流に証明(っぽいもの)をしてみる. 正直形式的べき級数をあまり良く理解しておらず, 雰囲気で書いているため不正確な記述になっているとは思う($\mathrm{mod}$ とか書くのも正確なのか知らないが $\mathrm{mod}\ x^n$ と書いた場合は多項式環のそれと同様に ($x^n$) で割ったものというノリで書いている).

Lemma: $x_j (j \ge 2^i)$ の係数が $0$ であるような形式的べき級数 $g_i$ が $g_i \equiv \exp(f)\ \mathrm{mod}\ x^{2^i}$ かつ ($g_i$ の定数項) $= 1$ を満たすなら $g_{i+1} \equiv \exp(f)\ \mathrm{mod}\ x^{2^{i+1}}$ かつ ($g_{i+1}$ の定数項) $= 1$ を満たす形式的べき級数 $g_{i+1}$ は $g_{i+1} \overset{\mathrm{def}}{=} \left( \left( g_i \cdot (f + 1 - \log g_i) \right) \ \mathrm{mod}\ x^{2^{i+1}} \right)$ と定めることで得られる.

これより初期値を $g_0 = 1$ と取り, $\lceil \log r \rceil$ 回反復を繰り返すことで解が得られる.

(Lemma の証明(っぽい何か))
まず $\left( g_i \cdot (f + 1 - \log g_i) \right) \ \mathrm{mod}\ x^{2^{i+1}}$ の定数項に注目すると ($g_i$ の定数項) $= 1$ および ($\log g_i$ の定数項) $= 0$ より ($g_{i+1}$ の定数項) は $1$ となることが分かる.
次に $a(x)$ を $g_i(x)$ の $x^j (1 \le j < 2^i)$ の項からなる形式的べき級数, $b(x)$ を $\exp(x)$ の $x^j (2^i \le j < 2^{i+1})$ の項からなる形式的べき級数, $c(x)$ を $\exp(x)$ の $x^j (j \ge 2^{i+1})$ の項からなる形式的べき級数とする. このとき $\exp(f) \equiv g_i\ \mathrm{mod}\ x^{2^i}$ から $\exp(f(x)) = 1 + a(x) + b(x) + c(x)$ が成り立つ.
よって $\exp(f) - g_{i+1} \equiv (1 + a + b) - (1 + a) \cdot (\log(1 + a + b + c) + 1 - \log (1 + a)) \overset{?}{\equiv} 0\ \mathrm{mod}\ x^{2^{i+1}}$ を示せば良い(式 (1)).
$\log (1 + a + b + c) = \sum_{i \ge 1} \frac{(-1)^{i+1}}{i}(a + b + c)^i$ であり $x^{2^{i+1}}$ 以降の項は考えなくて良いので $c$ の寄与は考えなくて良いのと $b$ の寄与についても $b^i (i \ge 2)$ については考える必要がない.
ゆえに $\log (1 + a + b + c) \equiv \sum_{i \ge 1} \frac{(-1)^{i+1}}{i}a^i + b \sum_{i \ge 0} (-1)^i a^i\ \mathrm{mod}\ x^{2^{i+1}}$ が成り立ち, $\log (1 + a + b + c) \equiv \log (1 + a) + \frac{b}{1 + a}\ \mathrm{mod}\ x^{2^{i+1}}$ と書きなおすことができる.
これを式 (1) に代入することで成り立つことが示せた.

(関数)
polynomial_exp$(a, r)$: 多項式 $f$ の係数 $a$($a[0] = 0$) を与えたとき, それに対応する形式的べき級数 $F$ とする. このとき $\exp (F)$ の $x^i (0 \le i < 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;
}

vector<int> differentiate(const vector<int>& a)
{
    const int n = (int)a.size();
    vector<int> res(n - 1);
    for(int i = 0; i < n - 1; ++i){
        res[i] = mul(a[i + 1], i + 1);
    }
    return res;
}

vector<int> integral(const vector<int>& a)
{
    const int n = (int)a.size();
    vector<int> res(n + 1, 0);
    for(int i = 0; i < n; ++i){
        res[i + 1] = mul(a[i], inverse(i + 1));
    }
    return res;
}

vector<int> polynomial_log(const vector<int>& a, int r)
{
    assert(a[0] == 1);
    if(r == 0) return {0};
    const int n = (int)a.size();
    const vector<int> num = differentiate(a);
    const vector<int> den = polynomial_inverse(a, r - 1);
    const vector<int> res = convolute(num, den, min(n, r) - 1, r - 1, r - 1);
    return integral(res);
}

vector<int> polynomial_exp(const vector<int>& a, int r)
{
    assert(a[0] == 0);
    const int n = (int)a.size();
    vector<int> ans = {1};
    for(int t = 1; t < r;){
        t <<= 1;
        vector<int> res = polynomial_log(ans, t);
        for(int i = 0; i < t; ++i) res[i] = sub((int)0, res[i]);
        res[0] = add(res[0], 1);
        for(int i = 0; i < min(n, t); ++i){
            res[i] = add(res[i], a[i]);
        }
        ans = convolute(ans, res, t / 2, t, t);
    }
    ans.resize(r);
    return ans;
}

verify 用の問題

yosupo さんの library checker : Exp of Formal Power Series 提出コード