ARC114 C - Sequence Scores
解法
解説の解法と違ったので書いておく。
数列を左から決めていったとき、番目の数の操作回数への寄与を考える。を選んだとき操作回数が増えないのは、となるが存在し、かつ、 であるようなときであり、これ以外の場合は操作回数が1増える。
よって、操作回数が増えない確率 は と計算できるので、これを使って計算できる。
コード
use crate::mod_int::{set_mod_int, ModInt}; fn main() { let (r, w) = (std::io::stdin(), std::io::stdout()); let mut sc = IO::new(r.lock(), w.lock()); set_mod_int(998244353); let n: usize = sc.read(); let m: usize = sc.read(); let inv_m = ModInt::from(1) / m; let mut skip_p = vec![ModInt::from(0); m]; let mut pow = vec![ModInt::from(1); m]; let mut ans = ModInt::from(0); for _ in 0..n { for i in 0..m { ans += (ModInt::from(1) - skip_p[i]) * inv_m; } for i in 0..m { skip_p[i] += inv_m * pow[i]; pow[i] *= ModInt::from(m - 1 - i) * inv_m; } } ans *= ModInt::from(m).pow(n); println!("{}", ans.value()); } pub mod mod_int { type ModInternalNum = i64; thread_local!( static MOD: std::cell::RefCell<ModInternalNum> = std::cell::RefCell::new(0); ); pub fn set_mod_int<T: ToInternalNum>(v: T) { MOD.with(|x| x.replace(v.to_internal_num())); } fn modulo() -> ModInternalNum { MOD.with(|x| *x.borrow()) } pub struct ModInt(ModInternalNum); impl Clone for ModInt { fn clone(&self) -> Self { Self(self.0) } } impl Copy for ModInt {} impl ModInt { fn internal_new<T>(v: T) -> Self where ModInternalNum: From<T>, { let mut v = ModInternalNum::from(v); let m = modulo(); if v >= m { v %= m; } Self(v) } pub fn internal_pow(&self, mut e: ModInternalNum) -> Self { let mut result = 1; let mut cur = self.0; let modulo = modulo(); while e > 0 { if e & 1 == 1 { result *= cur; result %= modulo; } e >>= 1; cur = (cur * cur) % modulo; } Self(result) } pub fn pow<T: ToInternalNum>(&self, e: T) -> Self { self.internal_pow(e.to_internal_num()) } pub fn value(&self) -> ModInternalNum { self.0 } } pub trait ToInternalNum { fn to_internal_num(&self) -> ModInternalNum; } impl ToInternalNum for ModInt { fn to_internal_num(&self) -> ModInternalNum { self.0 } } macro_rules! impl_primitive { ($primitive:ident) => { impl From<$primitive> for ModInt { fn from(v: $primitive) -> Self { let v = v as ModInternalNum; Self::internal_new(v) } } impl ToInternalNum for $primitive { fn to_internal_num(&self) -> ModInternalNum { *self as ModInternalNum } } }; } impl_primitive!(u8); impl_primitive!(u16); impl_primitive!(u32); impl_primitive!(u64); impl_primitive!(usize); impl_primitive!(i8); impl_primitive!(i16); impl_primitive!(i32); impl_primitive!(i64); impl_primitive!(isize); impl<T: ToInternalNum> std::ops::AddAssign<T> for ModInt { fn add_assign(&mut self, rhs: T) { let mut rhs = rhs.to_internal_num(); let m = modulo(); if rhs >= m { rhs %= m; } self.0 += rhs; if self.0 >= m { self.0 -= m; } } } impl<T: ToInternalNum> std::ops::Add<T> for ModInt { type Output = ModInt; fn add(self, rhs: T) -> Self::Output { let mut res = self; res += rhs; res } } impl<T: ToInternalNum> std::ops::SubAssign<T> for ModInt { fn sub_assign(&mut self, rhs: T) { let mut rhs = rhs.to_internal_num(); let m = modulo(); if rhs >= m { rhs %= m; } if rhs > 0 { self.0 += m - rhs; } if self.0 >= m { self.0 -= m; } } } impl<T: ToInternalNum> std::ops::Sub<T> for ModInt { type Output = Self; fn sub(self, rhs: T) -> Self::Output { let mut res = self; res -= rhs; res } } impl<T: ToInternalNum> std::ops::MulAssign<T> for ModInt { fn mul_assign(&mut self, rhs: T) { let mut rhs = rhs.to_internal_num(); let m = modulo(); if rhs >= m { rhs %= m; } self.0 *= rhs; self.0 %= m; } } impl<T: ToInternalNum> std::ops::Mul<T> for ModInt { type Output = Self; fn mul(self, rhs: T) -> Self::Output { let mut res = self; res *= rhs; res } } impl<T: ToInternalNum> std::ops::DivAssign<T> for ModInt { fn div_assign(&mut self, rhs: T) { let mut rhs = rhs.to_internal_num(); let m = modulo(); if rhs >= m { rhs %= m; } let inv = Self(rhs).internal_pow(m - 2); self.0 *= inv.value(); self.0 %= m; } } impl<T: ToInternalNum> std::ops::Div<T> for ModInt { type Output = Self; fn div(self, rhs: T) -> Self::Output { let mut res = self; res /= rhs; res } } } pub struct IO<R, W: std::io::Write>(R, std::io::BufWriter<W>); impl<R: std::io::Read, W: std::io::Write> IO<R, W> { pub fn new(r: R, w: W) -> IO<R, W> { IO(r, std::io::BufWriter::new(w)) } pub fn write<S: ToString>(&mut self, s: S) { use std::io::Write; self.1.write_all(s.to_string().as_bytes()).unwrap(); } pub fn read<T: std::str::FromStr>(&mut self) -> T { use std::io::Read; let buf = self .0 .by_ref() .bytes() .map(|b| b.unwrap()) .skip_while(|&b| b == b' ' || b == b'\n' || b == b'\r' || b == b'\t') .take_while(|&b| b != b' ' && b != b'\n' && b != b'\r' && b != b'\t') .collect::<Vec<_>>(); unsafe { std::str::from_utf8_unchecked(&buf) } .parse() .ok() .expect("Parse error.") } pub fn usize0(&mut self) -> usize { self.read::<usize>() - 1 } pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> { (0..n).map(|_| self.read()).collect() } pub fn chars(&mut self) -> Vec<char> { self.read::<String>().chars().collect() } }