AtCoder Beginner Contest 456 G

https://atcoder.jp/contests/abc456/tasks/abc456_g

久しぶりに解説を見てみたら、母関数を使う方法が書いてありました。
Sをxで分割して、それぞれの領域の長さを n_iとして、 n日のうち連続する休日が k日までの場合の数を L(n, k)とすると、求める場合の数は、

 \displaystyle \prod_{i}{L(n_i, k)}

となります。なので、この L(n, k)を求めればよいです。
休日をo、勤務日をxで表すと、oは k個まで続いていいので、x、xo、…、xo…o(oが k個)を組み合わせて結合すればいいです。ただしこれだと必ずxから始まるので n+1の長さを k+1までの長さで分割することにします。例えば長さ3を1以下で分割するなら、

1 1 1 1 => x x x x => xxx
2 1 1   => xo x x  => oxx
1 2 1   => x xo x  => xox
1 1 2   => x x xo  => xxo
2 2     => xo xo   => oxo

の5通りになります。
一般化のために n個を k個以内で分割することを考えて、これを母関数にすると、

 \displaystyle G_k(x) = \sum_{i=1}^{\infty}{(\sum_{j=1}^{k}{x^j})^i}

となって、この x^nの係数が求める場合の数となります。 i個に分割して、 i番目の長さが jです。
ここまでは考えていたのですが、内側の和が x^kで終わるのが怖すぎて、ここでストップしてしまいました。解説ではここからさらに計算しています。

 \displaystyle G_k(x) = \sum_{i=1}^{\infty}{(\frac{x-x^{k+1}}{1-x})^i}
 \displaystyle \phantom{G_k(x)} = \frac{x-x^{k+1}}{1-x}\frac{1}{1-\frac{x-x^{k+1}}{1-x}}
 \displaystyle \phantom{G_k(x)} = \frac{x-x^{k+1}}{1-2x+x^{k+1}}

こうしてみると kが大きいと速く計算できそうな気がします。だから、 1-2xでくくると、

 \displaystyle \phantom{G_k(x)} = \frac{x-x^{k+1}}{1-2x}\frac{1}{1+\frac{x^{k+1}}{1-2x}}
 \displaystyle \phantom{G_k(x)} = \frac{x-x^{k+1}}{1-2x}\sum_{i=0}^{\infty}{(\frac{-x^{k+1}}{1-2x})^i}
 \displaystyle \phantom{G_k(x)} = (x-x^{k+1})\sum_{i=0}^{\infty}{\frac{(-1)^ix^{(k+1)i}}{(1-2x)^i}}

ここで \frac{1}{(1-2x)^i}を考えます。その前に \frac{1}{(1-x)^i}を考えると、これは (1 + x + x^2 + \cdots)^iなので、 x^jの係数は i種類の中から j個を選ぶ重複組み合わせになります。 \frac{1}{(1-2x)^i}なら x 2xに置き換えるだけなので、

 \displaystyle \frac{1}{(1-2x)^i} = \sum_{j=0}^{\infty}{_jH_i2^jx^j}

だから個別の係数はすぐに計算できます。なので、 G_k(x) x^nの係数は、 O(\frac{n}{k+1})で計算できます。

// Count Holidays
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

// ax = by + 1 (a, b > 0)
fn linear_diophantine(a: i64, b: i64) -> Option<(i64, i64)> {
    if a == 1 {
        return Some((1, 0))
    }
    
    let q = b / a;
    let r = b % a;
    if r == 0 {
        return None
    }
    let (x1, y1) = linear_diophantine(r, a)?;
    Some((-q * x1 - y1, -x1))
}

fn inverse(a: i64, d: i64) -> i64 {
    let (x, _y) = linear_diophantine(a, d).unwrap();
    if x >= 0 {
        x % d
    }
    else {
        x % d + d
    }
}

fn pow(n: i64, e: u32, d: i64) -> i64 {
    if e == 0 {
        1
    }
    else if e == 1 {
        n
    }
    else if e % 2 == 1 {
        pow(n, e-1, d) * n % d
    }
    else {
        let m = pow(n, e/2, d);
        m * m % d
    }
}


//////////////////// HCalculator ////////////////////

const D: i64 = 998244353;

struct HCalculator {
    facs: Vec<i64>,
    inv_facs: Vec<i64>
}

impl HCalculator {
    fn calc(&self, n: i64, m: i64) -> i64 {
        self.facs[(n+m-1) as usize] * self.inv_facs[m as usize] % D
                                    * self.inv_facs[(n-1) as usize] % D
    }
    
    fn build(N: usize) -> HCalculator {
        let mut facs: Vec<i64> = vec![1; N+1];
        let mut inv_facs: Vec<i64> = vec![1; N+1];
        for n in 2..N+1 {
            facs[n] = facs[n-1] * (n as i64) % D;
            inv_facs[n] = inv_facs[n-1] * inverse(n as i64, D) % D
        }
        HCalculator { facs, inv_facs }
    }
}


//////////////////// process ////////////////////

fn read_input() -> String {
    let _N: i64 = read();
    let S: String = read();
    S
}

fn divide(S: String) -> Vec<i64> {
    let mut ns: Vec<i64> = vec![];
    let mut n: i64 = 0;
    for c in S.chars() {
        if c == '.' {
            n += 1
        }
        else {
            if n > 0 {
                ns.push(n)
            }
            n = 0
        }
    }
    if n > 0 {
        ns.push(n)
    }
    ns
}

use std::collections::BTreeMap;

fn frequency(v: &Vec<i64>) -> Vec<(i64, u32)> {
    let mut m: BTreeMap<i64, u32> = BTreeMap::new();
    for &n in v.iter() {
        let e = m.entry(n).or_insert(0);
        *e += 1
    }
    m.into_iter().rev().collect::<Vec<_>>()
}

// 2^e % D
fn pow2(e: i64) -> i64 {
    if e == 0 {
        1
    }
    else if e == 1 {
        2
    }
    else if e % 2 == 1 {
        pow2(e-1) * 2 % D
    }
    else {
        let n = pow2(e/2);
        n * n % D
    }
}

// nを1~kの整数で分割するときの場合の数
fn num_divs(n: i64, k: i64, hc: &HCalculator) -> i64 {
    let mut s: i64 = 0;
    for p in 0..(n-1)/(k+1)+1 {
        let sign: i64 = if p % 2 == 0 { 1 } else { -1 };
        // pのときの最小次数
        let q1 = (k+1)*p+1;
        let m1 = n - q1;
        s = (s + sign * pow2(m1) * hc.calc(p+1, m1)).rem_euclid(D);
        let q2 = (k+1)*(p+1);
        let m2 = n - q2;
        if m2 >= 0 {
            s = (s - sign * pow2(m2) * hc.calc(p+1, m2)).rem_euclid(D)
        }
    }
    s
}

fn F(S: String) {
    let N = S.len();
    let ns = divide(S);
    let freq = frequency(&ns);
    let hc = HCalculator::build((N*2) as usize);
    let mut accs: Vec<(Vec<i64>, u32)> = vec![];
    for (n, f) in freq {
        let acc: Vec<i64> = (0..n+1).map(|k| num_divs(n+1, k+1, &hc)).collect();
        accs.push((acc, f))
    }
    
    // i日連続以下の場合の数
    let mut a: Vec<i64> = vec![1; N+1];
    let mut full: i64 = 1;
    for i in 1..N+1 {
        a[i] = full;
        for (acc, f) in accs.iter() {
            if i >= acc.len() {
                break
            }
            a[i] = a[i] * pow(acc[i], *f, D) % D;
            if i == acc.len() - 1 {
                // ここで最後、あとは同じ値が続くとする
                full = full * pow(acc[i], *f, D) % D
            }
        }
    }
    
    for i in 1..N+1 {
        println!("{}", (a[i] - a[i-1]).rem_euclid(D))
    }
}

fn main() {
    let S = read_input();
    F(S)
}