AtCoder Beginner Contest 422 G

https://atcoder.jp/contests/abc422/tasks/abc422_g

1は簡単です。箱1のボールの個数がAの倍数の母関数は、
\displaystyle G_1(x)=1+x^A+x^{2A}+\cdots = \frac{1}{1-x^A}
なので、3つの箱全体での母関数は、
\displaystyle G(x) = G_1(x)G_2(x)G_3(x) = \frac{1}{(1-x^A)(1-x^B)(1-x^C)}
分母は精々8項なので、O(N)でxNの項まで計算できます。

2は1のようにはいきません。ボールを区別するので指数型の母関数にする必要があります。例えば5個のボールを箱1に3個、箱2に2個ボールを入れるとすると、5C2通りあります。これは、 \displaystyle (1+\frac{x^3}{3!})(1+\frac{x^2}{2!}) = 1+\frac{x^2}{2!}+\frac{x^3}{3!}+\frac{x^5}{2!3!}なのでx5の項は \frac{_5C_2}{5!}x^5となって、ボールを区別したときの場合の数が計算できます。
しかし、これはふつうに計算するとO(N2)なので無理です。
ただし、BCが大きいときは間に合います。A ≦ B ≦ Cとなるように順番を変えて、箱2と箱3の母関数を掛け算します。この計算量はO( (N/B)(N/C))で、箱1の母関数を掛けるときはxNの係数だけ分かればいいので、O(N)で計算できます。
BCが小さいときだけ別の計算方法を考えないといけません。(N/B)(N/C) ≤ 108とのきは上の方法で計算するとして、(N/B)(N/C) ＞ 108のとき、N ≤ 3×105を使うと、BC ＜ 9×102で、ABC ＜ 2.7×104となります。
もう一つの計算方法は、状態遷移を考えるものです。箱1のボールの個数がAで割っていくつになるかなどを状態とします。M=ABCとすると、状態数はMでふつうに状態遷移を行列で表すと行列の掛け算にO(M3)かかるので全然間に合いません。状態遷移を1回ずつ行っていくとO(NM)かかるので、NM ＜ 8.1×109で、各状態は3つの状態に遷移するので、間に合わない感じです（あとでやってみたら1774msでギリ間に合いました）。
もっと速い方法があります。遷移行列は各状態は3つの状態に遷移するので疎な行列になります。そうすると、Berlekamp-Massey法が使えます。これは、漸化式を高速に求める方法です。状態数はMなので、状態の個数はM次元のベクトルで表されますが、それがM回状態遷移すれば各回のベクトルで1次従属になるものが存在するので、M+1個以下のベクトルの漸化式で表されます。このベクトルの計算にO(M2)かかります。そして、ベクトルを適当な重みを付けてスカラー化して、Berlekamp-Massey法を使うと漸化式が求められます。確実とは言えないのですが（例えば、重みが全部1だと3倍3倍になって確実にうまくいきません）、これはベクトルの漸化式にもなります。この漸化式でN番目の項を求めればよいですが、これは高速に計算できます。漸化式を多項式と考えて、xNをその多項式割った余りを計算するのと同じになります。計算量はO(M2logN)で、これは1.3×1010くらいになります。試してみると速いものではこれで10倍くらい速いです。Nがもっと大きいとこの方法は非常に有効になります。
手元で最も時間がかかりそうなケースを試してみたのですが、全然返ってきません。問題はずいぶん楽なものしかないようです。

// Balls and Boxes
#![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()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}

// 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<const D: i64>(a: i64) -> i64 {
    let (x, _y) = linear_diophantine(a, D).unwrap();
    x.rem_euclid(D)
}


//////////////////// Polynomial ////////////////////

use std::collections::HashMap;

// 非ゼロの係数の項だけ残す
type Polynomial = Vec<(usize, i64)>;

fn mul(f1: &Polynomial, f2: &Polynomial) -> Polynomial {
    let mut dic: HashMap<usize, i64> = HashMap::new();
    for &(n1, c1) in f1.iter() {
        for &(n2, c2) in f2.iter() {
            let n = n1 + n2;
            let c = c1 * c2;
            let e = dic.entry(n).or_insert(0);
            *e += c
        }
    }
    
    let mut f: Vec<(usize, i64)> = dic.into_iter().
                                        filter(|&(_, c)| c != 0).collect();
    f.sort();
    f
}

// fの定数項は1である前提
fn inverse3(f: &Polynomial, N: usize) -> Polynomial {
    let mut v: Vec<i64> = vec![0; N+1];
    let mut w: Vec<i64> = vec![0; N+1];
    w[0] = 1;
    for i in 0..N+1 {
        let c1 = w[i];
        if c1 == 0 {
            continue
        }
        v[i] = c1;
        for &(j, c2) in f.iter() {
            if i + j > N {
                break
            }
            w[i+j] -= c1 * c2
        }
    }
    v.into_iter().enumerate().collect::<Polynomial>()
}


//////////////////// ModInt ////////////////////

use std::cmp::Ordering;
use std::ops::{Add, AddAssign, Sub, SubAssign,
               Mul, MulAssign, Div, DivAssign, Neg, Rem};

#[derive(Copy, Clone, Debug)]
struct ModInt<const D: i64>(i64);

impl<const D: i64> From<i64> for ModInt<D> {
    fn from(x: i64) -> Self {
        Self(x.rem_euclid(D))
    }
}
impl<const D: i64> PartialEq for ModInt<D> {
    fn eq(&self, other: &Self) -> bool {
        self.0 == other.0
    }
}

impl<const D: i64> PartialEq<i64> for ModInt<D> {
    fn eq(&self, other: &i64) -> bool {
        self.0 == *other
    }
}

impl<const D: i64> Eq for ModInt<D> {}

impl<const D: i64> PartialOrd for ModInt<D> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<const D: i64> Ord for ModInt<D> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.cmp(&other.0)
    }
}

impl<const D: i64> Add for ModInt<D> {
    type Output = Self;
    
    fn add(self, other: Self) -> Self {
        Self((self.0 + other.0) % D)
    }
}

impl<const D: i64> Add<i64> for ModInt<D> {
    type Output = Self;
    
    fn add(self, other: i64) -> Self {
        Self((self.0 + other) % D)
    }
}

impl<const D: i64> Add<ModInt<D>> for i64 {
    type Output = ModInt<D>;
    
    fn add(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>((self + other.0).rem_euclid(D))
    }
}

impl<const D: i64> AddAssign for ModInt<D> {
    fn add_assign(&mut self, other: ModInt<D>) {
        self.0 = (self.0 + other.0).rem_euclid(D)
    }
}

impl<const D: i64> AddAssign<i64> for ModInt<D> {
    fn add_assign(&mut self, other: i64) {
        self.0 = (self.0 + other).rem_euclid(D)
    }
}

impl<const D: i64> Sub for ModInt<D> {
    type Output = Self;
    
    fn sub(self, other: Self) -> Self {
        Self((self.0 - other.0).rem_euclid(D))
    }
}

impl<const D: i64> Sub<i64> for ModInt<D> {
    type Output = Self;
    
    fn sub(self, other: i64) -> Self {
        Self((self.0 - other).rem_euclid(D))
    }
}

impl<const D: i64> Sub<ModInt<D>> for i64 {
    type Output = ModInt<D>;
    
    fn sub(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>((self - other.0).rem_euclid(D))
    }
}

impl<const D: i64> SubAssign for ModInt<D> {
    fn sub_assign(&mut self, other: Self) {
        self.0 = (self.0 - other.0).rem_euclid(D)
    }
}

impl<const D: i64> SubAssign<i64> for ModInt<D> {
    fn sub_assign(&mut self, other: i64) {
        self.0 = (self.0 - other).rem_euclid(D)
    }
}

impl<const D: i64> Mul for ModInt<D> {
    type Output = Self;
    
    fn mul(self, other: Self) -> Self {
        Self(self.0 * other.0 % D)
    }
}

impl<const D: i64> Mul<i64> for ModInt<D> {
    type Output = Self;
    
    fn mul(self, other: i64) -> Self {
        Self(self.0 * other % D)
    }
}

impl<const D: i64> Mul<ModInt<D>> for i64 {
    type Output = ModInt<D>;
    
    fn mul(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>(self * other.0 % D)
    }
}

impl<const D: i64> MulAssign for ModInt<D> {
    fn mul_assign(&mut self, other: Self) {
        self.0 = self.0 * other.0 % D
    }
}

impl<const D: i64> MulAssign<i64> for ModInt<D> {
    fn mul_assign(&mut self, other: i64) {
        self.0 = self.0 * other % D
    }
}

impl<const D: i64> Div for ModInt<D> {
    type Output = Self;
    
    fn div(self, other: Self) -> Self {
        Self(self.0 * inverse::<D>(other.0) % D)
    }
}

impl<const D: i64> Div<i64> for ModInt<D> {
    type Output = Self;
    
    fn div(self, other: i64) -> Self {
        Self(self.0 * inverse::<D>(other) % D)
    }
}

impl<const D: i64> Div<ModInt<D>> for i64 {
    type Output = ModInt<D>;
    
    fn div(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>(self * inverse::<D>(other.0) % D)
    }
}

impl<const D: i64> DivAssign for ModInt<D> {
    fn div_assign(&mut self, other: Self) {
        self.0 = self.0 * inverse::<D>(other.0) % D
    }
}

impl<const D: i64> DivAssign<i64> for ModInt<D> {
    fn div_assign(&mut self, other: i64) {
        self.0 = self.0 * inverse::<D>(other) % D
    }
}

impl<const D: i64> Neg for ModInt<D> {
    type Output = Self;
    
    fn neg(self) -> Self::Output {
        Self(if self.0 == 0 { 0 } else { D - self.0 })
    }
}

impl<const D: i64> Rem for ModInt<D> {
    type Output = Self;
    
    fn rem(self, other: ModInt<D>) -> ModInt<D> {
        ModInt::<D>(self.0 % other.0)
    }
}

impl<const D: i64> Rem<i64> for ModInt<D> {
    type Output = Self;
    
    fn rem(self, other: i64) -> Self {
        Self(self.0 % other)
    }
}

use std::iter::Sum;

impl<const D: i64> Sum for ModInt<D> {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self::from(0), |acc, x| acc + x)
    }
}


//////////////////// PolyExp ////////////////////

// 非ゼロの係数の項だけ残す
type PolyExp = Vec<(usize, IntD)>;

fn mul2(f1: &PolyExp, f2: &PolyExp, fac: &Vec<IntD>) -> PolyExp {
    let N = fac.len() - 1;
    let mut dic: HashMap<usize, IntD> = HashMap::new();
    for &(n1, c1) in f1.iter() {
        for &(n2, c2) in f2.iter() {
            let n = n1 + n2;
            if n > N {
                break
            }
            let c = c1 * c2 * fac[n] / fac[n1] / fac[n2];
            let e = dic.entry(n).or_insert(IntD::from(0));
            *e += c
        }
    }
    
    let mut f: PolyExp = dic.into_iter().filter(|&(_, c)| c != 0).collect();
    f.sort();
    f
}


//////////////////// Berlekamp-Massey ////////////////////

fn BerlekampMassey<const D: i64>(a: &Vec<i64>) -> Vec<i64> {
    let N = a.len();
    let mut C: Vec<i64> = vec![1];  // 漸化式
    let mut prev_C = C.clone();
    let mut prev_d: i64 = 1;        // 前回のずれ
    let mut m: usize = 1;           // 何回続けて漸化式が合っているか
    for n in 0..N {
        let mut d: i64 = 0;
        for i in 0..C.len() {
            d = (d + C[i]*a[n-i]).rem_euclid(D);
        }
        if d == 0 {     // 漸化式が合っている
            m += 1
        }
        else {
            // 漸化式を今回と前回の結果で更新する
            let mut new_C = C.clone();
            for (i, &c) in prev_C.iter().enumerate() {
                while i + m >= new_C.len() {
                    new_C.push(0)
                }
                new_C[i+m] = (new_C[i+m] - c * d % D * inverse::<D>(prev_d))
                                                                .rem_euclid(D);
            }
            
            if new_C.len() > C.len() {
                prev_C = C.clone();
                prev_d = d;
                m = 1
            }
            else {
                m += 1
            }
            C = new_C
        }
    }
    C
}


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

fn read_input() -> (usize, usize, usize, usize) {
    let v: Vec<usize> = read_vec();
    let (N, A, B, C) = (v[0], v[1], v[2], v[3]);
    (N, A, B, C)
}

// 1 - x^n
fn create(n: usize) -> Polynomial {
    vec![(0, 1), (n, -1)]
}

fn F1(N: usize, A: usize, B: usize, C: usize) -> i64 {
    // G(x) = 1/(1-x^A)(1-x^B)(1-x^C)
    let inv_G = mul(&create(A), &mul(&create(B), &create(C)));
    let G = inverse3(&inv_G, N);
    G[N].1 % D
}

const D: i64 = 998244353;

type IntD = ModInt::<D>;

fn make_factorial_table(N: usize) -> Vec<IntD> {
    let mut fac: Vec<IntD> = vec![IntD::from(1); N+1];
    for n in 1..N+1 {
        fac[n] = fac[n-1] * (n as i64)
    }
    fac
}

// 1 + x^n/n! + x^2n/2n! + ...
fn create2(n: usize, N: usize) -> PolyExp {
    let mut f: PolyExp = PolyExp::new();
    for i in (0..N+1).step_by(n) {
        f.push((i, IntD::from(1)))
    }
    f
}

// BCが大きいときは母関数の計算をする
fn F2a(N: usize, A: usize, B: usize, C: usize) -> i64 {
    let fac = make_factorial_table(N);
    
    // G(x) = GA(x)GB(x)GC(x)
    // GA(x) = 1+x^A/A!+x^(2A)/(2A)!+...
    let GA = create2(A, N);
    let GB = create2(B, N);
    let GC = create2(C, N);
    let GBC = mul2(&GB, &GC, &fac);
    
    // GAを掛けるときx^Nの係数だけ計算する
    let L1 = GA.len();
    let L2 = GBC.len();
    let mut k: usize = 0;
    let mut l: usize = 0;
    let mut c: IntD = IntD::from(0);
    while k < L1 && l < L2 {
        let (n1, c1) = GA[k];
        let (n2, c2) = GBC[L2-l-1];
        if n1 + n2 == N {
            c += c1 * c2 * fac[N] / (fac[n1] * fac[n2]);
            k += 1;
            l += 1
        }
        else if n1 + n2 < N {
            k += 1
        }
        else {
            l += 1
        }
    }
    c.0
}

// 状態遷移してパスの数を数える
fn transition(v: &Vec<IntD>, A: usize, B: usize, C: usize) -> Vec<IntD> {
    let M = A * B * C;
    let mut w: Vec<IntD> = vec![IntD::from(0); M];
    for c in 0..C {
        for b in 0..B {
            for a in 0..A {
                let s = a + A * (b + c * B);
                if a < A - 1 {
                    w[s+1] += v[s]
                }
                else {
                    w[s+1-A] += v[s]
                }
                if b < B - 1 {
                    w[s+A] += v[s]
                }
                else {
                    w[s+A-A*B] += v[s]
                }
                if c < C - 1 {
                    w[s+A*B] += v[s]
                }
                else {
                    w[s+A*B-A*B*C] += v[s]
                }
            }
        }
    }
    w
}

fn initAndNearby(A: usize, B: usize, C: usize) -> Vec<Vec<IntD>> {
    let M = A * B * C;
    
    // 初期状態
    let mut v0: Vec<IntD> = vec![IntD::from(0); M];
    v0[0] = IntD::from(1);
    // 状態を保存しておく
    let mut vs: Vec<Vec<IntD>> = vec![vec![]; M*2];
    vs[0] = v0.clone();
    let mut v = v0.clone();
    for i in 1..M*2 {
        v = transition(&v, A, B, C);
        vs[i] = v.clone()
    }
    vs
}

fn find_coeffieces(A: usize, B: usize, C: usize,
                                vs: &Vec<Vec<IntD>>) -> Vec<IntD> {
    let M = A * B * C;
    if M == 1 {
        return vec![IntD::from(1), IntD::from(-3)]
    }
    
    // 乱数っぽい係数を作る
    let mut rs: Vec<i64> = vec![1; M];
    rs[1] = 2;
    for i in 2..M {
        rs[i] = (rs[i-1] + rs[i-2]) % D
    }
    
    // 状態をさっきの係数を使ってスカラーにする
    let mut a: Vec<IntD> = vec![IntD::from(0); M*2];
    for i in 0..M*2 {
        for j in 0..M {
            a[i] += vs[i][j] * rs[j]
        }
    }
    
    // 状態の漸化式を作る
    let b: Vec<i64> = a.into_iter().map(|n| n.0).collect();
    let c = BerlekampMassey::<D>(&b);
    c.into_iter().map(|n| IntD::from(n)).collect::<Vec<_>>()
}

// 商と余りを返す
fn div_IntD(f: &Vec<IntD>, g: &Vec<IntD>) -> (Vec<IntD>, Vec<IntD>) {
    let L1 = f.len();
    let L2 = g.len();
    if L1 < L2 {
        return (vec![IntD::from(0)], f.clone())
    }
    
    let L = L1 + 1 - L2;
    let mut r = f.clone();
    let mut q: Vec<IntD> = vec![IntD::from(0); L];
    for i in 0..L {
        let c = r[i] / g[0];
        for k in 0..L2 {
            r[i+k] -= c * g[k]
        }
        q[i] = c
    }
    (q, r[L..].to_vec())
}

fn mul_IntD(f: &Vec<IntD>, g: &Vec<IntD>) -> Vec<IntD> {
    let L1 = f.len();
    let L2 = g.len();
    let L = L1 + L2 - 1;
    let mut h: Vec<IntD> = vec![IntD::from(0); L];
    for i in 0..L1 {
        for j in 0..L2 {
            h[i+j] += f[i] * g[j]
        }
    }
    h
}

// x^N / fの余りを返す
fn pow(N: usize, f: &Vec<IntD>) -> Vec<IntD> {
    let L = f.len();
    if N < L - 1 {
        let mut g = vec![IntD::from(0); N+1];
        g[0] = IntD::from(1);
        g
    }
    else if N % 2 == 1 {
        let g = pow(N-1, f);
        let mut h: Vec<IntD> = vec![IntD::from(0); g.len()];
        for i in 0..g.len()-1 {
            h[i] = g[i+1]
        }
        for i in 0..h.len() {
            h[i] -= g[0] * f[i+1]
        }
        h
    }
    else {
        let g = pow(N/2, f);
        let g2 = mul_IntD(&g, &g);
        let (_, r) = div_IntD(&g2, f);
        r
    }
}

// BCが小さいときは状態遷移の計算をする
fn F2b(N: usize, A: usize, B: usize, C: usize) -> i64 {
    let vs = initAndNearby(A, B, C);
    if N < vs.len() {
        return vs[N][0].0
    }
    
    let cs = find_coeffieces(A, B, C, &vs);
    let f = pow(N, &cs);
    let L = f.len();
    let mut s = IntD::from(0);
    for i in 0..L {
        s += f[i] * vs[L-1-i][0]
    }
    s.0
}

fn F2(N: usize, mut A: usize, mut B: usize, mut C: usize) -> i64 {
    // A <= B <= Cにする
    let mut v: Vec<usize> = vec![A, B, C];
    v.sort();
    (A, B, C) = (v[0], v[1], v[2]);
    
    if (N/B)*(N/C) + N < 100_000_000 {
        F2a(N, A, B, C)
    }
    else {
        F2b(N, A, B, C)
    }
}

fn main() {
    let (N, A, B, C) = read_input();
    println!("{}", F1(N, A, B, C));
    println!("{}", F2(N, A, B, C))
}