「FPS 24 題」全部解く: N「硬貨 2」

解きます

解く

式は一瞬で立ちますね

$$f = \prod_{k=1}^N \frac{1-x^{k\left(A_k+1\right)}}{1-x^k}$$

maspyさんの記事に書いてあるテクが使えそうです.

$$\begin{align}
f &= \exp\sum_{k=1}^N\log\frac{1-x^{k\left(A_k+1\right)}}{1-x^k} \\
&= \exp\sum_{k=1}^N\log\left(1-x^{k\left(A_k+1\right)}\right)-\log\left(1-x^k\right) \\
&= \exp\int\left(\sum_{k=1}^N\frac{-k\left(A_k+1\right)x^{k\left(A_k+1\right)-1}}{1-x^{k\left(A_k+1\right)}}-\frac{-kx^{k-1}}{1-x^k}\right)\mathrm dx \\
\end{align}$$
分数の羃の肩を見ると, みんな$k$とかついているため, 項数は$k$に反比例して減っていきます. $N$次以下の項は全体で$O\left(N\log N\right)$個しかないと言え, これを愚直に全て足してシグマは計算できます.

あとは指数関数が計算できればよいですが, これは$\log f-g=0$みたいな方程式をニュートン法で解けばよいです:
$$\begin{align}
f_{n+1}&=f_n-\frac{\log f_n-g}{\left(\log f_n-g\right)'} \\
&=f_n-\frac{\log f_n-g}{1/f_n} \\
&=f_n\left(1-\log f_n+g\right) \\
\end{align}$$

use ac_library::*;
use proconio::*;

type Mint = ModInt998244353;

fn main() {
    input! {
        n: usize,
        a: [usize; n],
    }
    // シグマ
    let mut s = vec![Mint::raw(0); n.next_power_of_two() - 1];
    for k in 1..=n {
        let g = k * (a[k - 1] + 1);
        // 左側の分数
        if g - 1 < s.len() {
            for v in s[g - 1..].iter_mut().step_by(g) {
                *v -= g;
            }
        }
        // 右側の分数
        for v in s[k - 1..].iter_mut().step_by(k) {
            *v += k;
        }
    }
    // 積分する
    s.insert(0, Mint::raw(0));
    for (i, v) in s.iter_mut().enumerate().skip(1) {
        *v /= i;
    }
    // 指数関数をニュートン法で計算する
    // 答え
    let mut exp = vec![Mint::raw(1)];
    // expの逆数
    let mut inv_exp = vec![Mint::raw(1)];
    // 十分な精度になるまで反復
    while exp.len() < s.len() {
        // 仮でinv_expのニュートン法をすすめる
        {
            let m = convolution(&inv_exp, &inv_exp);
            let m = convolution(&m, &exp);
            inv_exp.extend((0..inv_exp.len() - 1).map(|_| Mint::raw(0)));
            for i in 0..inv_exp.len() {
                inv_exp[i] = inv_exp[i] * 2 - m[i];
            }
        }
        // 答えの微分 d/dx[exp s]
        let dexp = exp
            .iter()
            .enumerate()
            .skip(1)
            .map(|(i, v)| *v * i)
            .collect::<Vec<_>>();
        // (log f)'
        let mut p = convolution(&inv_exp, &dexp);
        // 積分する
        p.insert(0, Mint::raw(0));
        for (i, v) in p.iter_mut().enumerate().skip(1) {
            *v /= i;
        }
        // 1 - log f + s
        let mut q = vec![Mint::raw(0); exp.len() * 2];
        for (i, v) in q.iter_mut().enumerate() {
            *v = s[i] - p.get(i).unwrap_or(&Mint::raw(0));
        }
        q[0] += 1;
        exp = convolution(&exp, &q);
        // 精度が無い部分を切り捨て
        exp.drain(q.len()..);
        // expを変更したのでinv_expも再計算する
        {
            inv_exp.drain(exp.len() / 2..);
            let m = convolution(&inv_exp, &inv_exp);
            let m = convolution(&m, &exp);
            inv_exp.extend((0..inv_exp.len()).map(|_| Mint::raw(0)));
            for i in 0..inv_exp.len() {
                inv_exp[i] = inv_exp[i] * 2 - m[i];
            }
        }
    }
    println!("{}", exp[n]);
}

提出(472 ms):

終わり

デバッグ出力を消し忘れて1WA出しました
以上です