結果

問題 No.3391 Line up Dominoes
コンテスト
ユーザー northward
提出日時 2025-12-05 21:10:56
言語 Rust
(1.83.0 + proconio)
結果
AC  
実行時間 994 ms / 3,000 ms
コード長 8,394 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 13,505 ms
コンパイル使用メモリ 399,696 KB
実行使用メモリ 7,720 KB
最終ジャッジ日時 2025-12-05 21:11:32
合計ジャッジ時間 34,574 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 23
権限があれば一括ダウンロードができます

ソースコード

diff #
raw source code 

#![allow(non_snake_case, unused_must_use, unused_imports)]
use std::io::{self, prelude::*};
use std::{fmt::*, ops::*};

const MOD: u32 = 998244353;

fn main() {
    let (stdin, stdout) = (io::read_to_string(io::stdin()).unwrap(), io::stdout());
    let (mut stdin, mut buffer) = (stdin.split_whitespace(), io::BufWriter::new(stdout.lock()));

    macro_rules! input {
        ($t: tt, $n: expr) => {
            (0..$n).map(|_| input!($t)).collect::<Vec<_>>()
        };
        (Chars) => {
            input!(String).chars().collect::<Vec<_>>()
        };
        (Usize1) => {
            stdin.next().unwrap().parse::<usize>().unwrap() - 1
        };
        ($t: ty) => {
            stdin.next().unwrap().parse::<$t>().unwrap()
        };
    }

    let N = input!(usize);
    let M = input!(usize);
    let K = input!(usize);

    let mut A = input!(usize, N);
    A.sort();

    let mut left = vec![0; N];
    let mut right = vec![N - 1; N];

    for i in 0..N {
        // left
        if A[i] - A[0] > K {
            let mut ok = i;
            let mut ng = 0;

            while ok - ng > 1 {
                let mid = (ok + ng) / 2;

                if A[i] - A[mid] <= K {
                    ok = mid;
                } else {
                    ng = mid;
                }
            }

            left[i] = ok;
        }

        // right
        if A[N - 1] - A[i] > K {
            let mut ok = i;
            let mut ng = N - 1;

            while ng - ok > 1 {
                let mid = (ok + ng) / 2;

                if A[mid] - A[i] <= K {
                    ok = mid;
                } else {
                    ng = mid;
                }
            }

            right[i] = ok;
        }
    }

    let mut dp = vec![ModInt::<MOD>(1); N];

    for _ in 0..M - 1 {
        let mut dp_nxt = vec![ModInt::<MOD>(0); N];

        let cs = CumulativeSum::from(&dp);

        for i in 0..N {
            dp_nxt[i] = cs.sum(left[i]..=right[i]);
        }

        dp = dp_nxt;
    }

    let mut ans = ModInt::<MOD>(0);

    for v in dp {
        ans += v;
    }

    writeln!(buffer, "{}", ans);
}

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, Default)]
pub struct ModInt<const P: u32>(u32);

impl<const P: u32> ModInt<P> {
    pub fn from_raw(value: u32) -> Self {
        assert!(value < P);
        Self(value)
    }

    pub fn pow(&self, mut x: u32) -> Self {
        let mut a = *self;
        let mut r = Self::from_raw(1);

        while x > 0 {
            if x & 1 == 1 {
                r *= a;
            }

            a *= a;
            x >>= 1;
        }

        r
    }

    pub fn inv(&self) -> Self {
        self.pow(P - 2)
    }
}

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

impl<const P: u32> Sub for ModInt<P> {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self::Output {
        Self((P + self.0 - rhs.0) % P)
    }
}

impl<const P: u32> Mul for ModInt<P> {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self::Output {
        Self(((self.0 as u64 * rhs.0 as u64) % P as u64) as u32)
    }
}

impl<const P: u32> Div for ModInt<P> {
    type Output = Self;
    fn div(self, rhs: Self) -> Self::Output {
        self * rhs.inv()
    }
}

impl<const P: u32> AddAssign for ModInt<P> {
    fn add_assign(&mut self, rhs: Self) {
        self.0 += rhs.0;
        self.0 %= P;
    }
}

impl<const P: u32> SubAssign for ModInt<P> {
    fn sub_assign(&mut self, rhs: Self) {
        self.0 += P - rhs.0;
        self.0 %= P;
    }
}

impl<const P: u32> MulAssign for ModInt<P> {
    fn mul_assign(&mut self, rhs: Self) {
        *self = self.clone() * rhs;
    }
}

impl<const P: u32> DivAssign for ModInt<P> {
    fn div_assign(&mut self, rhs: Self) {
        *self *= rhs.inv()
    }
}

impl<const P: u32> Neg for ModInt<P> {
    type Output = Self;
    fn neg(self) -> Self::Output {
        Self((P - self.0) % P)
    }
}

impl<const P: u32> Display for ModInt<P> {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
        write!(f, "{}", self.0)
    }
}

macro_rules! impl_op_for_modint {
    ($($t: ty), *) => {
        $(
            impl<const P: u32> From<$t> for ModInt<P> {
                fn from(value: $t) -> Self {
                    Self((P as $t + value % P as $t) as u32 % P)
                }
            }

            impl<const P: u32> Add<$t> for ModInt<P> {
                type Output = Self;
                fn add(self, rhs: $t) -> Self::Output {
                    self + Self::from(rhs)
                }
            }

            impl<const P: u32> Add<ModInt<P>> for $t {
                type Output = ModInt<P>;
                fn add(self, rhs: ModInt<P>) -> Self::Output {
                    Self::Output::from(self) + rhs
                }
            }

            impl<const P: u32> Sub<$t> for ModInt<P> {
                type Output = Self;
                fn sub(self, rhs: $t) -> Self::Output {
                    self - Self::from(rhs)
                }
            }

            impl<const P: u32> Sub<ModInt<P>> for $t {
                type Output = ModInt<P>;
                fn sub(self, rhs: ModInt<P>) -> Self::Output {
                    Self::Output::from(self) - rhs
                }
            }

            impl<const P: u32> Mul<$t> for ModInt<P> {
                type Output = Self;
                fn mul(self, rhs: $t) -> Self::Output {
                    self * Self::from(rhs)
                }
            }

            impl<const P: u32> Mul<ModInt<P>> for $t {
                type Output = ModInt<P>;
                fn mul(self, rhs: ModInt<P>) -> Self::Output {
                    Self::Output::from(self) * rhs
                }
            }

            impl<const P: u32> Div<$t> for ModInt<P> {
                type Output = Self;
                fn div(self, rhs: $t) -> Self::Output {
                    self / Self::from(rhs)
                }
            }

            impl<const P: u32> Div<ModInt<P>> for $t {
                type Output = ModInt<P>;
                fn div(self, rhs: ModInt<P>) -> Self::Output {
                    Self::Output::from(self) / rhs
                }
            }

            impl<const P: u32> AddAssign<$t> for ModInt<P> {
                fn add_assign(&mut self, rhs: $t) {
                    *self += Self::from(rhs)
                }
            }

            impl<const P: u32> SubAssign<$t> for ModInt<P> {
                fn sub_assign(&mut self, rhs: $t) {
                    *self -= Self::from(rhs)
                }
            }

            impl<const P: u32> MulAssign<$t> for ModInt<P> {
                fn mul_assign(&mut self, rhs: $t) {
                    *self *= Self::from(rhs)
                }
            }

            impl<const P: u32> DivAssign<$t> for ModInt<P> {
                fn div_assign(&mut self, rhs: $t) {
                    *self /= Self::from(rhs)
                }
            }
        )*
    };
}

impl_op_for_modint!(usize, isize, u64, i64, u32, i32);

pub struct CumulativeSum<T> {
    size: usize,
    prefix_sum: Vec<T>,
}

impl<T: std::ops::Add<Output = T> + Default + Clone + Copy> CumulativeSum<T> {
    pub fn from(array: &[T]) -> Self {
        let size = array.len();
        let mut prefix_sum = vec![T::default(); size + 1];

        for i in 0..size {
            prefix_sum[i + 1] = prefix_sum[i] + array[i];
        }

        Self { size, prefix_sum }
    }

    pub fn prefix_sum(&self, i: usize) -> T {
        self.prefix_sum[i + 1]
    }
}

impl<T: std::ops::Add<Output = T> + std::ops::Sub<Output = T> + Default + Clone + Copy>
    CumulativeSum<T>
{
    pub fn sum<R: std::ops::RangeBounds<usize>>(&self, range: R) -> T {
        let left = match range.start_bound() {
            std::ops::Bound::Included(&l) => l,
            std::ops::Bound::Excluded(&l) => l + 1,
            std::ops::Bound::Unbounded => 0,
        };

        let right = match range.end_bound() {
            std::ops::Bound::Included(&r) => r,
            std::ops::Bound::Excluded(&r) => r - 1,
            std::ops::Bound::Unbounded => self.size - 1,
        };

        if left == 0 {
            return self.prefix_sum(right);
        } else {
            return self.prefix_sum(right) - self.prefix_sum(left - 1);
        }
    }
}
0