#838796 (Rust) No.2211 Frequency Table of GCD

提出ソース
結果

問題	No.2211 Frequency Table of GCD
コンテスト
ユーザー	sakikuroe
提出日時	2023-02-10 21:55:12
言語	Rust (1.83.0 + proconio)
結果	TLE
実行時間	-
コード長	28,163 bytes
コンパイル時間	11,897 ms
コンパイル使用メモリ	403,688 KB
実行使用メモリ	26,600 KB
最終ジャッジ日時	2024-07-07 16:09:23
合計ジャッジ時間	17,618 ms
ジャッジサーバーID （参考情報）	judge5 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 3
other	AC * 1 TLE * 1 -- * 24
権限があれば一括ダウンロードができます
ソースコード

raw source code
#[macro_export]
macro_rules! read_line {
    ($($xs: tt)*) => {
        let mut buf = String::new();
        std::io::stdin().read_line(&mut buf).unwrap();
        let mut iter = buf.split_whitespace();
        expand!(iter, $($xs)*);
    };
}

#[macro_export]
macro_rules! expand {
    ($iter: expr,) => {};
    ($iter: expr, mut $var: ident : $type: tt, $($xs: tt)*) => {
        let mut $var = value!($iter, $type);
        expand!($iter, $($xs)*);
    };
    ($iter: expr, $var: ident : $type: tt, $($xs: tt)*) => {
        let $var = value!($iter, $type);
        expand!($iter, $($xs)*);
    };
}

#[macro_export]
macro_rules! value {
    ($iter:expr, ($($type: tt),*)) => {
        ($(value!($iter, $type)),*)
    };
    ($iter: expr, [$type: tt; $len: expr]) => {
        (0..$len).map(|_| value!($iter, $type)).collect::<Vec<_>>()
    };
    ($iter: expr, Chars) => {
        value!($iter, String).unwrap().chars().collect::<Vec<_>>()
    };
    ($iter: expr, $type: ty) => {
        if let Some(v) = $iter.next() {
            v.parse::<$type>().ok()
        } else {
            None
        }
    };
}

use std::{
    collections::{HashMap, VecDeque},
    fmt, iter, ops,
};

const MOD: usize = 998244353; // 119 * (1 << 23) + 1
const RANK: usize = 23;
const PRIMITIVE_ROOT: usize = 3;

#[macro_export]
macro_rules! m {
    ($x:expr) => {
        ModInt::new((MOD as isize + $x as isize) as usize)
    };
}

#[macro_export]
macro_rules! fps {
    ($($x:expr),*) => (
        FPS::new(vec![$(m!($x)),*])
    );
    ($x:expr; $n:expr) => (
        FPS::new(vec![m!($x); $n])
    );
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub struct ModInt {
    value: usize,
}

impl ModInt {
    pub fn new(value: usize) -> ModInt {
        ModInt { value: value % MOD }
    }

    pub fn value(&self) -> usize {
        self.value
    }

    pub fn inverse(&self) -> ModInt {
        // (a, b) -> (x, y) s.t. a * x + b * y = gcd(a, b)
        fn extended_gcd(a: isize, b: isize) -> (isize, isize) {
            if (a, b) == (1, 0) {
                (1, 0)
            } else {
                let (x, y) = extended_gcd(b, a % b);
                (y, x - (a / b) * y)
            }
        }

        let (x, _y) = extended_gcd(self.value() as isize, MOD as isize);
        ModInt::new((MOD as isize + x) as usize)
    }

    // a^n
    pub fn pow(&self, mut n: usize) -> ModInt {
        let mut res = ModInt::new(1);
        let mut x = *self;
        while n > 0 {
            if n % 2 == 1 {
                res = res * x;
            }
            x = x * x;
            n /= 2;
        }

        res
    }
}

impl ops::Add for ModInt {
    type Output = ModInt;
    fn add(self, other: Self) -> Self {
        ModInt::new(self.value + other.value)
    }
}

impl ops::Sub for ModInt {
    type Output = ModInt;
    fn sub(self, other: Self) -> Self {
        ModInt::new(MOD + self.value - other.value)
    }
}

impl ops::Mul for ModInt {
    type Output = ModInt;
    fn mul(self, other: Self) -> Self {
        ModInt::new(self.value * other.value)
    }
}

impl ops::Div for ModInt {
    type Output = ModInt;
    fn div(self, other: Self) -> Self {
        self * other.inverse()
    }
}

impl ops::AddAssign for ModInt {
    fn add_assign(&mut self, other: Self) {
        *self = *self + other;
    }
}

impl ops::SubAssign for ModInt {
    fn sub_assign(&mut self, other: Self) {
        *self = *self - other;
    }
}

impl ops::MulAssign for ModInt {
    fn mul_assign(&mut self, other: Self) {
        *self = *self * other;
    }
}

impl ops::DivAssign for ModInt {
    fn div_assign(&mut self, other: Self) {
        *self = *self / other;
    }
}

impl fmt::Display for ModInt {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.value())
    }
}

fn binary_exponentiation<T>(mut a: T, mut n: usize, id: T, mut f: impl FnMut(T, T) -> T) -> T
where
    T: Copy,
{
    let mut ans = id;
    while n != 0 {
        if n % 2 == 1 {
            ans = f(ans, a);
        }
        a = f(a, a);
        n = n / 2;
    }
    ans
}

// Cipolla
pub fn sqrt_mod(y: usize, p: usize) -> Option<usize> {
    let y = y % p;
    if p == 2 {
        Some(y)
    } else if y == 0 {
        Some(0)
    } else if binary_exponentiation(y, p / 2, 1, |x, y| x * y % p) != 1 {
        None
    } else {
        let mut b = 0;
        loop {
            let sqgen = (b * b + p - y) % p;
            if binary_exponentiation(sqgen, p / 2, 1, |x, y| x * y % p) != 1 {
                return Some(
                    binary_exponentiation([b, 1], (p + 1) / 2, [1, 0], |x, y| {
                        [
                            (x[0] * y[0] + (x[1] * y[1]) % p * sqgen) % p,
                            (x[0] * y[1] + x[1] * y[0]) % p,
                        ]
                    })[0],
                );
            }
            b = b + 1;
        }
    }
}

pub struct FftCache {
    rate: Vec<ModInt>,
    irate: Vec<ModInt>,
}

impl FftCache {
    pub fn new() -> Self {
        let mut root = vec![ModInt::new(0); RANK + 1];
        let mut iroot = vec![ModInt::new(0); RANK + 1];
        let mut rate = vec![ModInt::new(0); RANK - 1];
        let mut irate = vec![ModInt::new(0); RANK - 1];

        root[RANK] = ModInt::new(PRIMITIVE_ROOT).pow((MOD - 1) >> RANK);
        iroot[RANK] = root[RANK].inverse();
        for i in (0..RANK).rev() {
            root[i] = root[i + 1] * root[i + 1];
            iroot[i] = iroot[i + 1] * iroot[i + 1];
        }

        {
            let mut prod = ModInt::new(1);
            let mut iprod = ModInt::new(1);
            for i in 0..RANK - 1 {
                rate[i] = root[i + 2] * prod;
                irate[i] = iroot[i + 2] * iprod;
                prod *= iroot[i + 2];
                iprod *= root[i + 2];
            }
        }

        FftCache { rate, irate }
    }
}

pub fn conv(a: &Vec<ModInt>, b: &Vec<ModInt>, cache: &FftCache) -> Vec<ModInt> {
    let ntt = |a: &mut Vec<ModInt>| {
        let n = a.len();
        let h = n.trailing_zeros();

        for len in 0..h {
            let p = 1 << (h - len - 1);
            let mut rot = ModInt::new(1);
            for (s, offset) in (0..1 << len).map(|s| s << (h - len)).enumerate() {
                let (al, ar) = a[offset..offset + 2 * p].split_at_mut(p);
                for (al, ar) in al.iter_mut().zip(ar.iter_mut()) {
                    let l = *al;
                    let r = *ar * rot;
                    *al = l + r;
                    *ar = l - r;
                }
                rot *= cache.rate[(!s).trailing_zeros() as usize];
            }
        }
    };

    let intt = |a: &mut Vec<ModInt>| {
        let n = a.len();
        let h = n.trailing_zeros();

        for len in (1..=h).rev() {
            let p = 1 << (h - len);
            let mut irot = ModInt::new(1);
            for (s, offset) in (0..1 << (len - 1)).map(|s| s << (h - len + 1)).enumerate() {
                let (al, ar) = a[offset..offset + 2 * p].split_at_mut(p);
                for (al, ar) in al.iter_mut().zip(ar.iter_mut()) {
                    let l = *al;
                    let r = *ar;
                    *al = l + r;
                    *ar = (l - r) * irot;
                }
                irot *= cache.irate[(!s).trailing_zeros() as usize];
            }
        }
    };

    if a.len() <= 2 {
        let mut res = vec![ModInt::new(0); a.len() + b.len() - 1];
        for i in 0..a.len() {
            for j in 0..b.len() {
                res[i + j] += a[i] * b[j];
            }
        }
        return res;
    } else if b.len() <= 2 || a.len() + b.len() <= 60 {
        let mut res = vec![ModInt::new(0); a.len() + b.len() - 1];
        for j in 0..b.len() {
            for i in 0..a.len() {
                res[i + j] += a[i] * b[j];
            }
        }
        return res;
    }

    let s = a.len() + b.len() - 1;
    let t = s.next_power_of_two();

    let (mut a, mut b) = (a.to_owned(), b.to_owned());
    a.resize(t, ModInt::new(0));
    ntt(&mut a);
    b.resize(t, ModInt::new(0));
    ntt(&mut b);
    a.iter_mut().zip(b.iter()).for_each(|(x, y)| *x = *x * *y);
    intt(&mut a);
    a.resize(s, ModInt::new(0));
    let t_inv = ModInt::new(t).inverse();
    a.iter_mut().for_each(|x| *x = *x * t_inv);
    a
}

#[derive(Clone, PartialEq, Eq)]
pub struct FPS {
    coefficient: Vec<ModInt>,
}

impl FPS {
    pub fn new(coefficient: Vec<ModInt>) -> Self {
        FPS { coefficient }
    }

    pub fn get(&self, deg: usize) -> Option<&ModInt> {
        self.coefficient.get(deg)
    }

    pub fn len(&self) -> usize {
        self.coefficient.len()
    }

    pub fn shrinked(&self, deg: usize) -> Self {
        FPS::new(self.coefficient.iter().copied().take(deg + 1).collect())
    }

    fn newton_by(deg: usize, init: ModInt, rec: impl Fn(FPS, usize) -> FPS) -> FPS {
        let mut res = fps! {init.value()};
        while res.len() < deg + 1 {
            let d = res.coefficient.len() * 2;
            res = rec(res, d - 1).shrinked(std::cmp::min(d - 1, deg));
        }

        res
    }

    pub fn inverse(&self, deg: usize) -> Self {
        let mut fps = self.clone();

        if fps.len() < deg + 1 {
            fps.coefficient.resize(deg + 1, ModInt::new(0));
        }

        let rec = |g: FPS, d| g.clone() * (fps! {2} - g * fps.shrinked(d));

        Self::newton_by(deg, fps.coefficient[0].inverse(), rec)
    }

    pub fn derivative(&self) -> Self {
        FPS::new(
            self.coefficient
                .iter()
                .enumerate()
                .skip(1)
                .map(|(i, &x)| ModInt::new(i) * x)
                .collect(),
        )
    }

    pub fn integral(&self) -> Self {
        FPS::new(
            vec![ModInt::new(0)]
                .into_iter()
                .chain(
                    self.coefficient
                        .iter()
                        .enumerate()
                        .map(|(i, &x)| x / ModInt::new(i + 1)),
                )
                .collect(),
        )
    }

    pub fn log(&self, deg: usize) -> Self {
        assert_eq!(
            *self.coefficient.get(0).unwrap_or(&ModInt::new(0)),
            ModInt::new(1)
        );
        (self.derivative().shrinked(deg) * self.inverse(deg))
            .shrinked(deg)
            .integral()
            .shrinked(deg)
    }

    pub fn exp(&self, deg: usize) -> Self {
        assert_eq!(
            *self.coefficient.get(0).unwrap_or(&ModInt::new(0)),
            ModInt::new(0)
        );

        let rec = |g: FPS, d| ((self.shrinked(d) + fps! {1} - g.log(d)) * g).shrinked(d);

        Self::newton_by(deg, ModInt::new(1), rec)
    }

    pub fn sqrt(&self, deg: usize) -> Option<Self> {
        let mut offset = 0;
        for i in 0..self.coefficient.len() {
            if self.coefficient[i] == m!(0) {
                offset += 1;
            } else {
                break;
            }
        }

        if offset == self.coefficient.len() {
            Some(fps!(0))
        } else if offset % 2 == 1 {
            None
        } else {
            let f = FPS::new(self.clone().coefficient.into_iter().skip(offset).collect());
            let inv = m!(2).inverse();
            let rec = |g: FPS, d| (g.clone() + f.shrinked(d) * g.inverse(d)) * fps! {inv.value()};

            if let Some(sqrt) = sqrt_mod(f.coefficient[0].value(), MOD) {
                Some(FPS::new(
                    vec![m!(0); offset / 2]
                        .into_iter()
                        .chain(
                            Self::newton_by(deg - offset / 2, m!(sqrt), rec)
                                .coefficient
                                .into_iter(),
                        )
                        .collect(),
                ))
            } else {
                None
            }
        }
    }

    // If:
    //     f = x^{k}mg ([x^{0}]g = 1)
    // Then:
    //     f^{n} = x^{kn}m^{n}exp(nlog(g))
    pub fn pow(&self, n: usize, max_deg: usize) -> FPS {
        if n == 0 {
            return fps! {1};
        }

        if self.coefficient.iter().all(|k| *k == ModInt::new(0)) {
            return fps! {0};
        }

        let mut k = 0;
        while let Some(&x) = self.get(k) {
            if x == ModInt::new(0) {
                k += 1;
            } else {
                break;
            }
        }

        if k as u128 * n as u128 >= max_deg as u128 + 1 {
            return FPS::new(vec![ModInt::new(0)]);
        }

        let m = *self.get(k).unwrap();
        let m_inv = m.inverse();
        let g = FPS::new(
            self.coefficient
                .iter()
                .cloned()
                .skip(k)
                .map(|x| x * m_inv)
                .collect(),
        );

        let mut res = g.log(max_deg + 1 - k * n);
        res *= FPS::new(vec![ModInt::new(n)]);
        res = res.exp(max_deg);
        res *= FPS::new(vec![m.pow(n)]);
        FPS::new(
            iter::repeat(ModInt::new(0))
                .take(k * n)
                .chain(res.coefficient.into_iter().take(max_deg + 1 - k * n))
                .collect(),
        )
    }

    /// Returns:
    ///     f(x + c)
    pub fn polynomial_taylor_shift(&self, c: ModInt) -> FPS {
        let mut fps = self.clone();
        let mut fact = ModInt::new(1);
        for i in 1..self.len() {
            fact *= ModInt::new(i);
            fps.coefficient[i] *= fact;
        }
        fps.coefficient.reverse();
        fps = fps * FPS::new(vec![ModInt::new(0), c]).exp(self.len());
        fps.coefficient.resize(self.len(), ModInt::new(0));
        fps.coefficient.reverse();
        let mut fact = ModInt::new(1);
        for i in 1..self.len() {
            fact *= ModInt::new(i);
            fps.coefficient[i] /= fact;
        }

        fps
    }

    /// Returns:
    ///     f(exp(x))
    pub fn composition_exp(&self, max_deg: usize) -> FPS {
        let mut ab = vec![];
        for i in 0..self.coefficient.len() {
            ab.push((self.coefficient[i], ModInt::new(i)));
        }

        sum_of_exp_bx(ab, max_deg)
    }

    pub fn to_fpss(&self) -> FPSS {
        FPSS::new(self.coefficient.iter().cloned().enumerate().collect())
    }

    pub fn set(&mut self, i: usize, x: ModInt) {
        if self.coefficient.len() <= i {
            self.coefficient.resize(i + 1, m!(0));
        }
        self.coefficient[i] = x;
    }
}

impl ops::Add for FPS {
    type Output = FPS;
    fn add(self, other: Self) -> Self {
        let len = std::cmp::max(self.len(), other.len());
        let mut res = self.coefficient.clone();
        res.resize(len, ModInt::new(0));
        res.iter_mut()
            .zip(other.coefficient.iter())
            .for_each(|(x, y)| *x += *y);
        FPS::new(res)
    }
}

impl ops::Sub for FPS {
    type Output = FPS;
    fn sub(self, other: Self) -> Self {
        let len = std::cmp::max(self.len(), other.len());
        let mut res = self.coefficient.clone();
        res.resize(len, ModInt::new(0));
        res.iter_mut()
            .zip(other.coefficient.iter())
            .for_each(|(x, y)| *x -= *y);
        FPS::new(res)
    }
}

impl ops::Mul for FPS {
    type Output = FPS;
    fn mul(self, other: Self) -> Self {
        let cache = FftCache::new();
        let coefficient = conv(&self.coefficient, &other.coefficient, &cache);
        FPS::new(coefficient)
    }
}

impl ops::AddAssign for FPS {
    fn add_assign(&mut self, other: Self) {
        *self = self.clone() + other;
    }
}

impl ops::SubAssign for FPS {
    fn sub_assign(&mut self, other: Self) {
        *self = self.clone() - other;
    }
}

impl ops::MulAssign for FPS {
    fn mul_assign(&mut self, other: Self) {
        *self = self.clone() * other;
    }
}

impl fmt::Display for FPS {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for i in 0..self.coefficient.len() {
            let _ = write!(f, "{}x^{} ", self.coefficient[i], i);
            if i + 1 != self.coefficient.len() {
                let _ = write!(f, "+ ");
            }
        }
        write!(f, "")
    }
}

impl Ord for FPS {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.partial_cmp(&other).unwrap()
    }
}

impl PartialOrd for FPS {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some((self.len()).cmp(&other.len()).reverse())
    }
}

pub fn product_of_polynomial_sequence(v: Vec<FPS>, max_deg: usize) -> FPS {
    let mut que = v.into_iter().collect::<VecDeque<_>>();
    que.push_back(FPS::new(vec![ModInt::new(1)]));

    while que.len() >= 2 {
        let fps1 = que.pop_front().unwrap();
        let fps2 = que.pop_front().unwrap();
        que.push_back(fps1 * fps2.shrinked(max_deg));
    }

    que[0].clone()
}

pub fn division_of_polynomials(mut f: FPS, g: FPS) -> (FPS, FPS) {
    if f.len() < g.len() {
        return (FPS::new(vec![]), f);
    }

    let mut rf = f.clone();
    let mut rg = g.clone();
    rf.coefficient.reverse();
    rg.coefficient.reverse();
    let deg = rf.len() - g.len() + 1;
    rf.coefficient.resize(deg, ModInt::new(0));
    rg.coefficient.resize(deg, ModInt::new(0));
    rg = rg.inverse(deg);
    let mut q = rf * rg;
    q.coefficient.resize(deg, ModInt::new(0));
    q.coefficient.reverse();
    let mut h = q.clone() * g;
    h.coefficient.resize(f.len(), ModInt::new(0));
    f -= h;
    while f.len() > 0 && f.coefficient[f.len() - 1] == ModInt::new(0) {
        f.coefficient.pop();
    }
    (q, f)
}

pub fn multipoint_evaluation(f: FPS, p: Vec<ModInt>) -> Vec<ModInt> {
    let m = p.len();

    let mut g = vec![FPS::new(vec![ModInt::new(1)]); 2 * m.next_power_of_two()];
    for i in 0..m {
        g[i + m.next_power_of_two()] = FPS::new(vec![ModInt::new(MOD - 1) * p[i], ModInt::new(1)]);
    }
    for i in (1..m.next_power_of_two()).rev() {
        g[i] = g[2 * i].clone() * g[2 * i + 1].clone();
    }

    let mut h = vec![FPS::new(vec![ModInt::new(1)]); 2 * m.next_power_of_two()];
    h[1] = division_of_polynomials(f.clone(), g[1].clone()).1;

    for i in 2..2 * m.next_power_of_two() {
        h[i] = division_of_polynomials(h[i / 2].clone(), g[i].clone()).1;
    }

    h[m.next_power_of_two()..]
        .iter()
        .take(m)
        .map(|x| *x.coefficient.get(0).unwrap_or(&ModInt::new(0)))
        .collect()
}

pub fn polynomial_interpolation(x: Vec<ModInt>, y: Vec<ModInt>) -> FPS {
    assert_eq!(x.len(), y.len());
    let n = x.len();

    let l = product_of_polynomial_sequence(
        x.iter()
            .cloned()
            .map(|x| FPS::new(vec![ModInt::new(MOD - 1) * x, ModInt::new(1)]))
            .collect(),
        x.len() - 1,
    );
    let dl = FPS::new(
        (1..l.len())
            .map(|i| l.coefficient[i] * ModInt::new(i))
            .collect(),
    );

    let mut g = vec![FPS::new(vec![ModInt::new(1)]); 2 * n.next_power_of_two()];
    for i in 0..n {
        g[i + n.next_power_of_two()] = FPS::new(vec![ModInt::new(MOD - 1) * x[i], ModInt::new(1)]);
    }
    for i in (1..n.next_power_of_two()).rev() {
        g[i] = g[2 * i].clone() * g[2 * i + 1].clone();
    }

    let w = multipoint_evaluation(dl, x.clone());

    let mut a = y
        .iter()
        .zip(w.iter())
        .map(|(x, y)| *x / *y)
        .collect::<Vec<_>>();
    a.resize(n.next_power_of_two(), ModInt::new(0));

    let mut flr = vec![FPS::new(vec![ModInt::new(0)]); 2 * n.next_power_of_two()];
    for i in 0..n {
        flr[n.next_power_of_two() + i] = FPS::new(vec![a[i]]);
    }
    for i in (1..n.next_power_of_two()).rev() {
        flr[i] =
            flr[2 * i].clone() * g[2 * i + 1].clone() + flr[2 * i + 1].clone() * g[2 * i].clone();
    }
    flr[1].coefficient.resize(n, ModInt::new(0));

    flr[1].clone()
}

/// fracs := \{(a_{i}, b_{i}\}_{i}
/// Returns:
///     - fracs := \{(a_{i}, b_{i}\}_{i}
///     \sum_{i} \frac{a_{i}}{b_{i}}
fn sum_of_rationals(fracs: Vec<(FPS, FPS)>) -> (FPS, FPS) {
    let mut fracs = fracs.into_iter().collect::<VecDeque<_>>();
    while fracs.len() > 1 {
        let a = fracs.pop_front().unwrap();
        let b = fracs.pop_front().unwrap();
        fracs.push_back((a.0 * b.1.clone() + a.1.clone() * b.0, a.1 * b.1));
    }
    fracs[0].clone()
}

/// Returns:
///     - ab := \{(a_{i}, b_{i}\}_{i}
///     \sum_{i} a_{i} exp(b_{i}x)
pub fn sum_of_exp_bx(ab: Vec<(ModInt, ModInt)>, max_deg: usize) -> FPS {
    let mut fracs = vec![];
    for (a, b) in ab {
        fracs.push((
            FPS::new(vec![a]),
            FPS::new(vec![ModInt::new(1), ModInt::new(MOD - 1) * b]),
        ));
    }
    let (mut f, mut g) = sum_of_rationals(fracs);
    g = g.shrinked(max_deg + 1);
    f = f * g.inverse(max_deg + 1);
    f = f.shrinked(max_deg + 1);
    let mut fact_inv = (1..=max_deg)
        .fold(ModInt::new(1), |x, y| x * ModInt::new(y))
        .inverse();
    for i in (1..=max_deg).rev() {
        f.coefficient[i] *= fact_inv;
        fact_inv *= ModInt::new(i);
    }
    f
}

pub struct BinomicalCoefficient {
    factorial: Vec<ModInt>,
    factorial_inv: Vec<ModInt>,
}

impl BinomicalCoefficient {
    pub fn new(max_size: usize) -> BinomicalCoefficient {
        let mut factorial = vec![m!(1)];
        for i in 1..=max_size {
            factorial.push(factorial[i - 1] * ModInt::new(i));
        }

        let mut factorial_inv = vec![m!(0); max_size + 1];
        factorial_inv[max_size] = factorial[max_size].inverse();
        for i in (0..max_size).rev() {
            factorial_inv[i] = factorial_inv[i + 1] * ModInt::new(i + 1);
        }

        BinomicalCoefficient {
            factorial,
            factorial_inv,
        }
    }

    // nCr
    pub fn combination(&self, n: usize, r: usize) -> ModInt {
        if n < r {
            return ModInt::new(0);
        }
        self.factorial[n] * self.factorial_inv[r] * self.factorial_inv[n - r]
    }
}

/// Returns:
///     (1 - rx)^{-d}
pub fn binomial_theorem_negative(r: ModInt, d: usize, deg: usize) -> FPS {
    let mut res = vec![];
    let mut r_pow = m!(1);
    let bino = BinomicalCoefficient::new(deg + d - 1);
    for n in 0..=deg {
        res.push(bino.combination(n + d - 1, d - 1) * r_pow);
        r_pow *= r;
    }
    FPS::new(res)
}

/// Bostan-Mori
pub fn kth_term_of_linearly_recurrent_sequence(a: Vec<ModInt>, c: Vec<ModInt>) -> (FPS, FPS) {
    assert_eq!(a.len(), c.len());
    assert!(c.len() > 0);
    let p = FPS::new(a);
    let dim = p.len();
    let q = FPS::new(
        vec![ModInt::new(1)]
            .into_iter()
            .chain(c.into_iter().map(|x| x * ModInt::new(MOD - 1)))
            .collect::<Vec<_>>(),
    );
    ((p * q.clone()).shrinked(dim), q)
}

#[derive(Clone, PartialEq, Eq)]
pub struct FPSS {
    coefficient: HashMap<usize, ModInt>,
}

impl FPSS {
    pub fn new(coefficient: HashMap<usize, ModInt>) -> Self {
        FPSS { coefficient }
    }

    pub fn get(&self, deg: usize) -> Option<&ModInt> {
        self.coefficient.get(&deg)
    }

    pub fn len(&self) -> usize {
        self.coefficient.len()
    }

    pub fn shrinked(&self, deg: usize) -> Self {
        FPSS::new(
            self.coefficient
                .clone()
                .into_iter()
                .filter(|(k, _v)| *k <= deg)
                .collect(),
        )
    }

    pub fn derivative(&self) -> FPSS {
        FPSS::new(
            self.coefficient
                .clone()
                .into_iter()
                .filter(|(k, _v)| *k >= 1)
                .map(|(k, v)| (k - 1, m!(k) * v))
                .collect(),
        )
    }

    pub fn inverse(&self, deg: usize) -> FPS {
        div_sparse(fps! {1}, self.clone(), deg)
    }

    pub fn integral(&self) -> FPSS {
        FPSS::new(
            self.coefficient
                .clone()
                .into_iter()
                .map(|(k, v)| (k + 1, v / m!(k + 1)))
                .collect(),
        )
    }

    pub fn log(&self, deg: usize) -> FPS {
        assert_eq!(*self.coefficient.get(&0).unwrap_or(&m!(0)), m!(1));
        (div_sparse(self.clone().derivative().to_fps(), self.clone(), deg))
            .shrinked(deg)
            .integral()
            .shrinked(deg)
    }

    pub fn to_fps(&self) -> FPS {
        let mut res = vec![m!(0); self.coefficient.iter().map(|(k, _v)| k).max().unwrap_or(&0) + 1];
        for (k, v) in self.coefficient.iter() {
            res[*k] = *v;
        }
        FPS::new(res)
    }

    pub fn set(&mut self, i: usize, x: ModInt) {
        self.coefficient.insert(i, x);
    }
}

impl ops::Add for FPSS {
    type Output = FPSS;
    fn add(self, other: Self) -> Self {
        let mut res = self.clone();
        for (key, v) in other.coefficient.iter() {
            *res.coefficient.entry(*key).or_insert(m!(0)) += *v;
        }
        res
    }
}

impl ops::Sub for FPSS {
    type Output = FPSS;
    fn sub(self, other: Self) -> Self {
        let mut res = self.clone();
        for (key, v) in other.coefficient.iter() {
            *res.coefficient.entry(*key).or_insert(m!(0)) -= *v;
        }
        res
    }
}

use std::mem;

pub fn gcd(mut a: usize, mut b: usize) -> usize {
    if a < b {
        mem::swap(&mut a, &mut b);
    }
    while b != 0 {
        let temp = a % b;
        a = b;
        b = temp;
    }
    a
}

pub fn lcm(a: usize, b: usize) -> usize {
    a * b / gcd(a, b)
}

impl ops::Mul for FPSS {
    type Output = FPSS;
    fn mul(self, other: Self) -> Self {
        let mut res = HashMap::new();

        for (key1, m1) in self.coefficient.iter() {
            for (key2, m2) in other.coefficient.iter() {
                *res.entry(gcd(*key1, *key2)).or_insert(m!(0)) += *m1 * *m2;
            }
        }

        FPSS::new(res)
    }
}

impl ops::AddAssign for FPSS {
    fn add_assign(&mut self, other: Self) {
        *self = self.clone() + other;
    }
}

impl ops::SubAssign for FPSS {
    fn sub_assign(&mut self, other: Self) {
        *self = self.clone() - other;
    }
}

impl ops::MulAssign for FPSS {
    fn mul_assign(&mut self, other: Self) {
        *self = self.clone() * other;
    }
}

pub fn product_of_polynomial_sequence_fpss(mut v: Vec<FPSS>, deg: usize) -> FPS {
    v.sort_by_key(|fpss| *fpss.coefficient.iter().map(|x| x.0).max().unwrap());
    v.reverse();
    let mut res = fps! {1}.to_fpss();

    for fps in v {
        res = res * fps;
        res = res.shrinked(deg);
    }

    res.to_fps()
}

/// Returns:
///     f / g
pub fn div_sparse(mut f: FPS, mut g: FPSS, deg: usize) -> FPS {
    assert_ne!(g.coefficient.get(&0).unwrap_or(&m!(0)).value(), 0);
    let inv = g.coefficient.get(&0).unwrap().inverse();
    f.coefficient.iter_mut().for_each(|x| *x *= inv);
    g.coefficient.iter_mut().for_each(|x| *x.1 *= inv);
    let mut res = f.coefficient;
    res.resize(deg + 1, m!(0));
    for i in 0..=deg {
        for (deg, v) in g
            .coefficient
            .iter()
            .filter(|(deg, _v)| **deg >= 1 && i >= **deg)
        {
            res[i] = res[i] - *v * res[i - *deg]
        }
    }
    FPS::new(res)
}

/// Returns:
///     [x^{k}] f * g
pub fn mul_kth(f: FPSS, g: FPSS, k: usize) -> ModInt {
    let mut res = m!(0);
    for (k1, v1) in f.coefficient.into_iter().filter(|(k1, _v)| k >= *k1) {
        res += v1 * *g.coefficient.get(&(k - k1)).unwrap_or(&m!(0))
    }
    res
}

pub fn mul_fps_fpss(f: FPS, g: FPSS) -> FPS {
    let mut res = fps! {};
    for (k, x) in g.coefficient.iter() {
        res += FPS::new(
            iter::repeat(m!(0))
                .take(*k)
                .chain(f.coefficient.iter().map(|y| *y * *x))
                .collect(),
        );
    }
    res
}

fn main() {
    read_line!(n: usize, m: usize,);
    let n = n.unwrap();
    let m = m.unwrap();
    read_line!(a: [usize; n],);
    let a = a.into_iter().map(|x| x.unwrap()).collect::<Vec<_>>();
    let mut v = vec![];
    for i in 0..n {
        let mut f = fps! {1}.to_fpss();
        f.set(a[i], m!(1));
        v.push(f);
    }

    let p = product_of_polynomial_sequence_fpss(v, m);
    for i in 1..=m {
        println!("{}", p.get(i).unwrap_or(&m!(0)));
    }
}
yukicoder

結果

ソースコード
