AtCoder Beginner Contest 358 E

https://atcoder.jp/contests/abc358/tasks/abc358_e

母関数を考えればよいですが、組合せじゃなくて順序も考えなければならないので指数関数型の母関数ですね。

例えば、Aが2個、Bが1個で3文字を考えると、AAB、ABA、BAAの3通りが考えられます。これは母関数

のの係数に対応します。係数はですが、これに3!を掛けると、3になります。これはと書けばわかりやすくて、AABは3つあるから順列は3!だけど、Aが2つあって入れ替えても同じだから2!で割っています。

// Alphabet Tiles
#![allow(non_snake_case)]

use std::cmp::min;

const D: i64 = 998244353;


//////////////////// 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(a: i64, d: i64) -> i64 {
    let (x, _y) = linear_diophantine(a, d).unwrap();
    if x >= 0 {
        x % d
    }
    else {
        x % d + d
    }
}


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

use std::ops::Mul;

struct Polynomial {
    a: Vec<i64>,
    K: usize
}

impl Mul for Polynomial {
    type Output = Self;
    
    fn mul(self, other: Self) -> Self {
        let L1 = self.len();
        let L2 = other.len();
        let L = min(L1+L2-1, self.K+1);
        let mut a: Vec<i64> = vec![0; L];
        for i in 0..L1 {
            for j in 0..min(L2, L-i) {
                a[i+j] = (a[i+j] + self.a[i] * other.a[j]) % D
            }
        }
        Self { a, K: self.K }
    }
}

impl Polynomial {
    fn len(&self) -> usize {
        self.a.len()
    }
    
    fn apply(&self, x: i64) -> i64 {
        let v = self.a.iter().rev().fold(0, |y, &c| (y*x+c)%D);
        (v + D) % D
    }
}


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

fn read_input() -> (usize, Vec<usize>) {
    let K: usize = read();
    let C: Vec<usize> = read_vec();
    (K, C)
}

fn create_polynomial(c: usize, K: usize) -> Polynomial {
    let n = min(c, K);
    let mut a: Vec<i64> = vec![1; n+1];
    for i in 1..n+1 {
        a[i] = a[i-1] * inverse(i as i64, D) % D
    }
    
    Polynomial { a, K }
}

fn create_generating_function(K: usize, C: Vec<usize>) -> Polynomial {
    let mut G = create_polynomial(C[0], K);
    for c in C[1..].iter() {
        G = G * create_polynomial(*c, K)
    }
    G
}

fn F(K: usize, C: Vec<usize>) -> i64 {
    let G = create_generating_function(K, C);
    let mut s: i64 = 0;
    let mut fac: i64 = 1;
    for i in 1..G.len() {
        fac = fac * (i as i64) % D;
        s = (s + G.a[i] * fac) % D
    }
    s
}

fn main() {
    let (K, C) = read_input();
    println!("{}", F(K, C))
}