#1075622 (Rust) No.3118 Increment or Multiply

• 問題: No.3118 Increment or Multiply
• ユーザー: rhoo
• 提出日時: 2025-04-18 23:22:49
• 言語: Rust (1.83.0 + proconio)
• 結果: AC
• 実行時間: 22 ms / 2,000 ms
• コード長: 9,900 bytes
• コンパイル時間: 13,226 ms
• コンパイル使用メモリ: 378,592 KB
• 実行使用メモリ: 6,272 KB
• 最終ジャッジ日時: 2025-04-18 23:23:12
• 合計ジャッジ時間: 15,458 ms
• ジャッジサーバーID （参考情報）: judge3 / judge1

ファイルパターン	結果
sample	AC * 1
other	AC * 35

ソースコード

#![allow(unused_imports,non_snake_case,dead_code)]
use std::{cmp::Reverse as Rev,collections::*,iter::*};
use proconio::{marker::*,*};
// use rand::prelude::*;


#[fastout]
fn main(){
    input!{
        qs:[(usize,usize)],
    }

    // let mut rng=StdRng::seed_from_u64(0);
    // let mut qs=(0..1000).map(|_|(rng.gen_range(1..=100),rng.gen_range(1..=100))).collect::<Vec<_>>();

    let M=|a|M::new(a);
    let mut path=vec![];
    
    for &(n,a) in &qs{
        if a==1{
            println!("{}",M(n)*(n-1)/2);
            continue;
        }
        
        path.clear();
        path.push(n);

        let mut ans=M(0);
        let mut last=n;

        while last!=0{
            if last%a!=0{
                last=last/a*a;
            } else{
                last/=a;
            }
            path.push(last);
        }
        
        for w in path.windows(2){
            if w[0]%a!=0{
                ans+=M(w[0]-1+w[1])*(w[0]-w[1])/2;
            } else{
                let len=w[0]-w[1]-1;
                ans+=M(len)*(len+1)/2;
                ans+=w[1];
            }
        }

        // let verify=solve(n,a);
        // assert!(verify==ans);

        // if verify!=ans{
        //     eprintln!("{} {}",n,a);
        //     eprintln!("verify = {verify}");
        //     eprintln!("ans = {ans}");
        // }

        println!("{ans}");
    }
}


fn solve(n:usize,a:usize)->M{
    let mut dp=vec![usize::MAX/64;n+1];
    dp[n]=0;
    let mut sum=M::new(0);
    for i in (1..=n).rev(){
        sum+=M::new(dp[i]);
        dp[i-1]=dp[i-1].min(dp[i]+1);
        if i%a==0{
            dp[i/a]=dp[i/a].min(dp[i]+1);
        }
    }
    sum
}


type M = StaticModInt<998244353>;
type ModInt1000000007 = StaticModInt<1000000007>;



#[derive(Clone, Copy, PartialEq, Eq, Default, Hash)]
struct StaticModInt<const MODULO: u32> {
    val: u32,
}

impl<const MODULO: u32> ModIntBase for StaticModInt<MODULO> {
    fn modulus() -> u32 {
        MODULO
    }
    unsafe fn raw(val: u32) -> Self {
        StaticModInt { val }
    }
    fn val(self) -> u32 {
        self.val
    }
    fn inv(self) -> Self {
        assert!(self.val != 0, "divisor by zero");
        let inv_val = mod_inv_by_ext_gcd(self.val, MODULO);
        unsafe { Self::raw(inv_val) }
    }
}



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

impl<const MODULO: u32> std::fmt::Debug for StaticModInt<MODULO> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.val)
    }
}


macro_rules! impl_from_int {
    ($($t:ty),*) => {
        $(impl<const MODULO: u32> From<$t> for StaticModInt<MODULO> {
            fn from(x: $t) -> Self {
                StaticModInt::new(x)
            }
        })*
    }
}

impl_from_int!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);



impl<const MODULO: u32> std::ops::Add for StaticModInt<MODULO> {
    type Output = Self;
    fn add(self, rhs: Self) -> Self::Output {
        let mut sum = self.val + rhs.val;
        if sum >= MODULO {
            sum -= MODULO;
        }
        unsafe { Self::raw(sum) }
    }
}

impl<const MODULO: u32> std::ops::Sub for StaticModInt<MODULO> {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self::Output {
        let diff = if self.val >= rhs.val {
            self.val - rhs.val
        } else {
            MODULO + self.val - rhs.val
        };
        unsafe { Self::raw(diff) }
    }
}

impl<const MODULO: u32> std::ops::Mul for StaticModInt<MODULO> {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self::Output {
        let prod = (self.val as u64 * rhs.val as u64 % MODULO as u64) as u32;
        unsafe { Self::raw(prod) }
    }
}

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



impl<const MODULO: u32> std::ops::AddAssign for StaticModInt<MODULO> {
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl<const MODULO: u32> std::ops::SubAssign for StaticModInt<MODULO> {
    fn sub_assign(&mut self, rhs: Self) {
        *self = *self - rhs;
    }
}

impl<const MODULO: u32> std::ops::MulAssign for StaticModInt<MODULO> {
    fn mul_assign(&mut self, rhs: Self) {
        *self = *self * rhs;
    }
}

impl<const MODULO: u32> std::ops::DivAssign for StaticModInt<MODULO> {
    fn div_assign(&mut self, rhs: Self) {
        *self = *self / rhs;
    }
}



impl<const MODULO: u32> std::ops::Neg for StaticModInt<MODULO> {
    type Output = Self;
    fn neg(self) -> Self::Output {
        if self.val == 0 {
            self
        } else {
            unsafe { Self::raw(MODULO - self.val) }
        }
    }
}



impl<const MODULO: u32> std::str::FromStr for StaticModInt<MODULO> {
    type Err = <u32 as std::str::FromStr>::Err;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let x = s.parse::<u64>()?;
        Ok(StaticModInt::new(x))
    }
}



macro_rules! impl_mod_int_ops_sub {
    ($trait:ident, $func:ident, $assign_trait:ident, $assign_func:ident, $op:tt, ($($t:ty),*)) => {
        $(
            impl<const MODULO: u32> std::ops::$trait<$t> for StaticModInt<MODULO> {
                type Output = Self;
                fn $func(self, rhs: $t) -> Self::Output {
                    self $op StaticModInt::new(rhs)
                }
            }
            impl<const MODULO: u32> std::ops::$trait<StaticModInt<MODULO>> for $t {
                type Output = StaticModInt<MODULO>;
                fn $func(self, rhs: StaticModInt<MODULO>) -> Self::Output {
                    StaticModInt::new(self) $op rhs
                }
            }
            impl<const MODULO: u32> std::ops::$assign_trait<$t> for StaticModInt<MODULO> {
                fn $assign_func(&mut self, rhs: $t) {
                    *self = *self $op StaticModInt::new(rhs);
                }
            }
        )*
    }
}

macro_rules! impl_mod_int_ops {
    ($($t:tt)*) => {
        impl_mod_int_ops_sub!($($t)*, (i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize));
    }
}

impl_mod_int_ops!(Add, add, AddAssign, add_assign, +);
impl_mod_int_ops!(Sub, sub, SubAssign, sub_assign, -);
impl_mod_int_ops!(Mul, mul, MulAssign, mul_assign, *);
impl_mod_int_ops!(Div, div, DivAssign, div_assign, /);



impl<const MODULO: u32> std::iter::Sum for StaticModInt<MODULO> {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(StaticModInt::zero(), |acc, x| acc + x)
    }
}

impl<const MODULO: u32> std::iter::Product for StaticModInt<MODULO> {
    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(StaticModInt::one(), |acc, x| acc * x)
    }
}



trait RemEuclidU32 {
    fn rem_euclid_u32(self, modulus: u32) -> u32;
}

macro_rules! impl_rem_euclid_u32_for_small_signed {
    ($($ty:tt),*) => {
        $(
            impl RemEuclidU32 for $ty {
                fn rem_euclid_u32(self, modulus: u32) -> u32 {
                    (self as i64).rem_euclid(i64::from(modulus)) as _
                }
            }
        )*
    }
}

impl_rem_euclid_u32_for_small_signed!(i8, i16, i32, i64, isize);

impl RemEuclidU32 for i128 {
    fn rem_euclid_u32(self, modulus: u32) -> u32 {
        self.rem_euclid(i128::from(modulus)) as _
    }
}

macro_rules! impl_rem_euclid_u32_for_small_unsigned {
    ($($ty:tt),*) => {
        $(
            impl RemEuclidU32 for $ty {
                fn rem_euclid_u32(self, modulus: u32) -> u32 {
                    self as u32 % modulus
                }
            }
        )*
    }
}

macro_rules! impl_rem_euclid_u32_for_large_unsigned {
    ($($ty:tt),*) => {
        $(
            impl RemEuclidU32 for $ty {
                fn rem_euclid_u32(self, modulus: u32) -> u32 {
                    (self % (modulus as $ty)) as _
                }
            }
        )*
    }
}

impl_rem_euclid_u32_for_small_unsigned!(u8, u16, u32);
impl_rem_euclid_u32_for_large_unsigned!(u64, u128);


#[cfg(target_pointer_width = "32")]
impl_rem_euclid_u32_for_small_unsigned!(usize);

#[cfg(target_pointer_width = "64")]
impl_rem_euclid_u32_for_large_unsigned!(usize);



trait ModIntBase: Default
    + std::str::FromStr
    + From<i8>
    + From<i16>
    + From<i32>
    + From<i64>
    + From<i128>
    + From<isize>
    + From<u8>
    + From<u16>
    + From<u32>
    + From<u64>
    + From<u128>
    + From<usize>
    + Copy
    + Eq
    + std::hash::Hash
    + std::fmt::Display
    + std::fmt::Debug
    + std::ops::Neg<Output = Self>
    + std::ops::Add<Output = Self>
    + std::ops::Sub<Output = Self>
    + std::ops::Mul<Output = Self>
    + std::ops::Div<Output = Self>
    + std::ops::AddAssign
    + std::ops::SubAssign
    + std::ops::MulAssign
    + std::ops::DivAssign
{
    fn modulus() -> u32;
    unsafe fn raw(val: u32) -> Self;
    fn val(self) -> u32;
    fn inv(self) -> Self;

    fn new(val: impl RemEuclidU32) -> Self {
        unsafe { Self::raw(val.rem_euclid_u32(Self::modulus())) }
    }

    fn zero() -> Self {
        unsafe { Self::raw(0) }
    }

    fn one() -> Self {
        unsafe { Self::raw(1) }
    }

    fn pow(self, mut n: u64) -> Self {
        let mut x = self;
        let mut r = Self::one();
        while n > 0 {
            if n & 1 == 1 {
                r *= x;
            }
            x *= x;
            n >>= 1;
        }
        r
    }
}



fn mod_inv_by_ext_gcd(x: u32, modulus: u32) -> u32 {
    let (mut a, mut b) = (x as i64, modulus as i64);
    let (mut u, mut v) = (1i64, 0i64);
    while b != 0 {
        let t = a / b;

        a -= t * b;
        (a,b)=(b,a);

        u -= t * v;
        (u,v)=(v,u);
    }
    assert!(a==1);

    ((u % modulus as i64 + modulus as i64) % modulus as i64) as u32
}