AtCoder Beginner Contest 431 F

https://atcoder.jp/contests/abc431/tasks/abc431_f

例えば、Dが1でAをソートした列が1 2 3 4だとすると、1と2の有効列は

1 2
2 1

の2つが考えられますが、これに3を追加するとすると、どちらでも2の左か最後に追加するしかないので、どちらの場合も2通り追加できます。そうすると、4通りのなります。

4を追加するのも同様なので、結局8通りになります。
ただし、Aに同じ数があると少しややしくなります。入力例1だと、ソートして、

1 2 2 5

同じ数はまとめて数だけ数えて、

(1, 1), (2, 2), (5, 1)

となります。まず1に2を2個追加するとき、1があるから3つのマスから2を置くマスを2つ選ぶので、₃C₂となります。

// Almost Sorted 2
#![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)
}


//////////////////// CombinationCalculator ////////////////////

struct CombinationCalculator<const D: i64> {
    facts: Vec<i64>
}

impl<const D: i64> CombinationCalculator<D> {
    // nCm
    fn apply(&self, n: i64, m: i64) -> i64 {
        let num = self.facts[n as usize];
        let den = self.facts[m as usize] * self.facts[(n-m) as usize] % D;
        num * inverse::<D>(den) % D
    }
    
    fn create(N: usize) -> CombinationCalculator::<D> {
        let mut facts: Vec<i64> = vec![1; N+1];
        for n in 1..N+1 {
            facts[n] = facts[n-1] * (n as i64) % D
        }
        CombinationCalculator::<D> { facts }
    }
}


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

fn read_input() -> (i32, Vec<i32>) {
    let v: Vec<i32> = read_vec();
    let D = v[1];
    let A: Vec<i32> = read_vec();
    (D, A)
}


fn groupby(A: &Vec<i32>) -> Vec<(i32, usize)> {
    let mut v: Vec<(i32, usize)> = vec![];
    let mut prev = 0;
    let mut len = 0;
    for &a in A.iter() {
        if len == 0  || a == prev {
            len += 1
        }
        else {
            v.push((prev, len));
            len = 1
        }
        prev = a
    }
    v.push((prev, len));
    v
}

// indexがどこまでさかのぼれるかをVecにする
fn calc_back(v: &Vec<(i32, usize)>, D: i32) -> Vec<usize> {
    let mut w: Vec<usize> = vec![0; v.len()];
    let mut k: usize = 0;
    for i in 0..v.len() {
        loop {
            let (a, _) = v[i];
            if a <= v[k].0 + D {
                w[i] = k;
                break
            }
            k += 1
        }
    }
    w
}

fn accumulate(v: &Vec<(i32, usize)>) -> Vec<usize> {
    let mut a: Vec<usize> = vec![0; v.len()+1];
    for i in 0..v.len() {
        a[i+1] = a[i] + v[i].1
    }
    a
}

const E: usize = 998244353;

fn F(D: i32, mut A: Vec<i32>) -> usize {
    let N = A.len();
    let cc = CombinationCalculator::<{E as i64}>::create(N);
    A.sort();
    let v = groupby(&A);
    let w = calc_back(&v, D);
    let a = accumulate(&v);
    let mut n: usize = 1;
    for i in 0..v.len() {
        let k = w[i];
        let C = cc.apply((a[i+1] - a[k]) as i64, v[i].1 as i64) as usize;
        n = n * C % E
    }
    n
}

fn main() {
    let (D, A) = read_input();
    println!("{}", F(D, A))
}