#1139122 (Rust) No.3396 ChRisTmas memory

• 問題: No.3396 ChRisTmas memory
• ユーザー: rsk0315
• 提出日時: 2025-12-07 15:28:46
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 1,500 ms / 4,000 ms
• コード長: 6,886 bytes
• コンパイル時間: 22,999 ms
• コンパイル使用メモリ: 402,212 KB
• 実行使用メモリ: 7,852 KB
• 最終ジャッジ日時: 2025-12-07 15:29:27
• 合計ジャッジ時間: 29,002 ms
• ジャッジサーバーID （参考情報）: judge3 / judge4

ファイルパターン	結果
sample	AC * 2
other	AC * 40

ソースコード

use std::io::BufRead;

use nekolib::math::{Gcd, GcdRecip};
use proconio::{
    fastout, input,
    source::{Readable, Source},
};

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum Query {
    Q1(u64, u64),
    Q2(usize),
    Q3(u64),
}
use Query::*;

impl Readable for Query {
    type Output = Query;
    fn read<R: BufRead, S: Source<R>>(source: &mut S) -> Self::Output {
        match u32::read(source) {
            1 => {
                input! {
                    from source,
                    m: u64,
                    r: u64,
                }
                Q1(m, r)
            }
            2 => {
                input! {
                    from source,
                    k: usize,
                }
                Q2(k)
            }
            3 => {
                input! {
                    from source,
                    m: u64,
                }
                Q3(m)
            }
            _ => unreachable!(),
        }
    }
}

#[fastout]
fn main() {
    input! {
        query: [Query],
    }

    let mut a = vec![(0, 1, true)];
    for &q in &query {
        match q {
            Q1(m, r) => {
                let new = || {
                    let lj = a.iter().map(|&(_, mj, _)| mj).fold(1, |x, y| x * y % m);
                    let g = lj.gcd(m);
                    let mi = m / g;

                    let (si, _) = a.iter().fold((0, 1), |(sj, lj), (yj, mj, _)| ((sj + yj * lj) % m, lj * mj % m));

                    let num = r + m * g - si;
                    (num % g == 0).then(|| (num / g * (lj / g).gcd_recip(mi).1 % mi, mi))
                };

                if let Some((yi, mi)) = new() {
                    a.push((yi, mi, true));
                } else {
                    a.push((0, 0, false));
                }
            }
            Q2(k) => {
                for _ in 0..k {
                    a.pop();
                }
            }
            Q3(m) => {
                if let Some((res, _)) =
                    a.iter().try_fold((0, 1), |(sj, lj), (yj, mj, ok)| ok.then(|| ((sj + yj * lj) % m, lj * mj % m)))
                {
                    println!("{res}");
                } else {
                    println!("-1");
                }
            }
        }
    }
}

/// This module is bundled automatically.
/// See <https://rsk0315.github.io/nekolib/nekolib_doc/index.html> for documentation.
/// Commit: f384433ce1b82eac51a72c3142b4d046158d1e1d-dirty
#[allow(unused)]
#[allow(private_interfaces)]
pub mod nekolib {
    pub mod math {
        pub mod gcd {
            pub trait Gcd {
                fn gcd(self, other: Self) -> Self;
            }
            macro_rules! impl_uint {
                            ( $($ty:ty)* ) => { $(
                                impl Gcd for $ty {
                                    fn gcd(mut self, mut other: Self) -> Self {
                                        while other != 0 {
                                            let tmp = self % other;
                                            self = std::mem::replace(&mut other, tmp);
                                        }
                                        self
                                    }
                                }
                            )* }
                        }
            impl_uint! { u8 u16 u32 u64 u128 usize }
        }
        #[allow(unused_imports)]
        pub use gcd::*;
        pub mod gcd_recip {
            pub trait GcdRecip: Sized {
                fn gcd_recip(self, other: Self) -> (Self, Self);
            }
            macro_rules! impl_uint {
                            ($t:ty) => {
                                impl GcdRecip for $t {
                                    fn gcd_recip(self, other: Self) -> (Self, Self) {
                                        assert!(other > 0);
                                        let a = self % other;
                                        if a == 0 {
                                            return (other, 0);
                                        }

                                        let mut s = other;
                                        let mut t = a;
                                        let mut m0 = 0;
                                        let mut m1 = 1;
                                        while t > 0 {
                                            let u = s / t;
                                            s -= t * u;
                                            // m0 -= m1 * u;
                                            let v = (m1 * u) % other;
                                            m0 = if m0 < v { m0 + other - v } else { m0 - v };
                                            std::mem::swap(&mut s, &mut t);
                                            std::mem::swap(&mut m0, &mut m1);
                                        }
                                        (s, m0 % (other / s))
                                    }
                                }
                            };
                            ( $($t:ty)* ) => { $(impl_uint!($t);)* };
                        }
            macro_rules! impl_int {
                            ($t:ty) => {
                                impl GcdRecip for $t {
                                    fn gcd_recip(self, other: Self) -> (Self, Self) {
                                        assert!(other > 0);
                                        let a = self.rem_euclid(other);
                                        if a == 0 {
                                            return (other, 0);
                                        }

                                        let mut s = other;
                                        let mut t = a;
                                        let mut m0 = 0;
                                        let mut m1 = 1;
                                        while t > 0 {
                                            let u = s / t;
                                            s -= t * u;
                                            m0 -= m1 * u;
                                            std::mem::swap(&mut s, &mut t);
                                            std::mem::swap(&mut m0, &mut m1);
                                        }
                                        if m0 < 0 {
                                            m0 += other / s;
                                        }
                                        (s, m0)
                                    }
                                }
                            };
                            ( $($t:ty)* ) => { $(impl_int!($t);)* };
                        }
            impl_uint!(u8 u16 u32 u64 u128 usize);
            impl_int!(i8 i16 i32 i64 i128 isize);
        }
        #[allow(unused_imports)]
        pub use gcd_recip::*;
    }
}