結果
問題 | No.2587 Random Walk on Tree |
ユーザー | akakimidori |
提出日時 | 2023-11-10 20:48:29 |
言語 | Rust (1.77.0 + proconio) |
結果 |
AC
|
実行時間 | 1,629 ms / 10,000 ms |
コード長 | 21,376 bytes |
コンパイル時間 | 15,209 ms |
コンパイル使用メモリ | 391,984 KB |
実行使用メモリ | 36,212 KB |
最終ジャッジ日時 | 2024-09-27 12:54:54 |
合計ジャッジ時間 | 48,342 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge1 |
(要ログイン)
テストケース
テストケース表示入力 | 結果 | 実行時間 実行使用メモリ |
---|---|---|
testcase_00 | AC | 1 ms
6,816 KB |
testcase_01 | AC | 1 ms
6,816 KB |
testcase_02 | AC | 1 ms
6,816 KB |
testcase_03 | AC | 1 ms
6,812 KB |
testcase_04 | AC | 2 ms
6,944 KB |
testcase_05 | AC | 2 ms
6,944 KB |
testcase_06 | AC | 2 ms
6,944 KB |
testcase_07 | AC | 1 ms
6,944 KB |
testcase_08 | AC | 2 ms
6,944 KB |
testcase_09 | AC | 1 ms
6,944 KB |
testcase_10 | AC | 2 ms
6,940 KB |
testcase_11 | AC | 4 ms
6,940 KB |
testcase_12 | AC | 19 ms
6,944 KB |
testcase_13 | AC | 27 ms
6,944 KB |
testcase_14 | AC | 7 ms
6,940 KB |
testcase_15 | AC | 1,309 ms
16,736 KB |
testcase_16 | AC | 766 ms
12,428 KB |
testcase_17 | AC | 834 ms
11,952 KB |
testcase_18 | AC | 127 ms
6,940 KB |
testcase_19 | AC | 1,484 ms
26,108 KB |
testcase_20 | AC | 1,408 ms
18,032 KB |
testcase_21 | AC | 1,485 ms
20,468 KB |
testcase_22 | AC | 1,463 ms
36,212 KB |
testcase_23 | AC | 1,454 ms
23,796 KB |
testcase_24 | AC | 1,542 ms
18,648 KB |
testcase_25 | AC | 1,149 ms
30,096 KB |
testcase_26 | AC | 1,517 ms
19,696 KB |
testcase_27 | AC | 1,456 ms
19,832 KB |
testcase_28 | AC | 1,629 ms
18,424 KB |
testcase_29 | AC | 1,555 ms
17,520 KB |
testcase_30 | AC | 1,587 ms
17,652 KB |
testcase_31 | AC | 1,585 ms
17,528 KB |
testcase_32 | AC | 1,528 ms
17,140 KB |
testcase_33 | AC | 1 ms
6,940 KB |
testcase_34 | AC | 1,105 ms
23,044 KB |
testcase_35 | AC | 1,130 ms
23,052 KB |
testcase_36 | AC | 1,133 ms
22,920 KB |
testcase_37 | AC | 1,160 ms
23,688 KB |
testcase_38 | AC | 1,491 ms
17,488 KB |
testcase_39 | AC | 1,507 ms
17,596 KB |
ソースコード
fn main() { type M = ModInt<998244353>; input! { n: usize, m: usize, s: usize1, t: usize1, e: [(usize1, usize1); n - 1], } let mut g = vec![vec![]; n]; for (a, b) in e { g[a].push(b); g[b].push(a); } let root = s; let mut topo = vec![root]; let mut parent = vec![n; n]; for i in 0..n { let v = topo[i]; for u in g[v].clone() { g[u].retain(|p| *p != v); parent[u] = v; topo.push(u); } } let mut size = vec![1i32; n]; for &v in topo.iter().rev() { g[v].sort_by_key(|u| -size[*u]); size[v] += g[v].iter().map(|u| size[*u]).sum::<i32>(); } let solve_path = |a: Vec<Vec<Vec<M>>>| -> Vec<Vec<M>> { let a = a .into_iter() .map(|a| { let mut b = vec![vec![vec![]; 2]; 2]; b[0][1] = a[0].clone(); b[1][0] = a[0].clone(); b[1][1] = a[1].clone(); b }) .collect::<Vec<_>>(); let ans = product( a, |l, r| { let mut res = vec![vec![vec![]; 2]; 2]; for (a, l) in l.iter().enumerate() { for (b, l) in l.iter().enumerate() { for (c, r) in r.iter().enumerate() { for (d, r) in r.iter().enumerate() { if b + c == 1 { continue; } let v = l.convolution(r); if (c, b) == (0, 0) { res[a][d].sub_assign(&v); } else { res[a][d].add_assign(&v); } } } } } res }, |a| a[1][1].len(), ); vec![ans[0][1].clone(), ans[1][1].clone()] }; let child_product = |a: Vec<Vec<Vec<M>>>| -> Vec<Vec<M>> { product( a, |a, b| { let mut c = vec![vec![]; 2]; for (i, a) in a.iter().enumerate() { for (j, b) in b.iter().enumerate() { if (i, j) != (0, 0) { c[i & j].add_assign(&a.convolution(b)); } } } c }, |a| a[1].len(), ) }; let lift = |a: Vec<Vec<M>>| -> Vec<Vec<M>> { let c = [[M::zero(), -M::one()], [M::one(), -M::one()]]; let mut res = vec![vec![]; 2]; res[0].add_assign(&c[0].convolution(&a[1])); res[1].add_assign(&c[1].convolution(&a[1])); res[1].sub_assign(&c[0].convolution(&a[0])); res }; let calc = recurse(|rec, mut v: usize| -> Vec<Vec<M>> { let mut poly = vec![]; loop { let mut a = g[v].iter().skip(1).map(|u| rec(*u)).collect::<Vec<_>>(); a.push(vec![vec![], vec![M::one()]]); let a = child_product(a); poly.push(lift(a)); if let Some(u) = g[v].get(0) { v = *u; } else { break; } } solve_path(poly) }); let mut nu = vec![]; let mut de = vec![]; let mut pos = t; let mut ban = n; let mut geta = 0; loop { let mut a = g[pos] .iter() .filter(|p| **p != ban) .map(|u| calc(*u)) .collect::<Vec<_>>(); a.push(vec![vec![], vec![M::one()]]); let a = child_product(a); nu.push(a[1].clone()); de.push(lift(a)); if pos == s { break; } ban = pos; pos = parent[pos]; geta += 1; } let mut nu = product(nu, |a, b| a.convolution(&b), |a| a.len()); nu.splice(0..0, (0..geta).map(|_| M::zero())); let mut de = solve_path(de)[1].clone(); let mut k = m; while k > 0 { let mut f = de.clone(); for f in f[1..].iter_mut().step_by(2) { *f = -*f; } nu = nu .convolution(&f) .into_iter() .skip(k & 1) .step_by(2) .collect(); de = de.convolution(&f).into_iter().step_by(2).collect(); k >>= 1; } println!("{}", nu[0]); } fn product<T, F, G>(mut a: Vec<T>, mul: F, size: G) -> T where F: Copy + Fn(T, T) -> T, G: Copy + Fn(&T) -> usize, { assert!(!a.is_empty()); if a.len() == 1 { return a.pop().unwrap(); } let sum = a.iter().map(size).sum::<usize>(); let mut c = 0; for i in 0..a.len() { if c >= sum - c || i == a.len() - 1 { let r = a.split_off(i); return mul(product(a, mul, size), product(r, mul, size)); } c += size(&a[i]); } unreachable!() } // ---------- begin input macro ---------- // reference: https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8 #[macro_export] macro_rules! input { (source = $s:expr, $($r:tt)*) => { let mut iter = $s.split_whitespace(); input_inner!{iter, $($r)*} }; ($($r:tt)*) => { let s = { use std::io::Read; let mut s = String::new(); std::io::stdin().read_to_string(&mut s).unwrap(); s }; let mut iter = s.split_whitespace(); input_inner!{iter, $($r)*} }; } #[macro_export] macro_rules! input_inner { ($iter:expr) => {}; ($iter:expr, ) => {}; ($iter:expr, $var:ident : $t:tt $($r:tt)*) => { let $var = read_value!($iter, $t); input_inner!{$iter $($r)*} }; } #[macro_export] macro_rules! read_value { ($iter:expr, ( $($t:tt),* )) => { ( $(read_value!($iter, $t)),* ) }; ($iter:expr, [ $t:tt ; $len:expr ]) => { (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>() }; ($iter:expr, chars) => { read_value!($iter, String).chars().collect::<Vec<char>>() }; ($iter:expr, bytes) => { read_value!($iter, String).bytes().collect::<Vec<u8>>() }; ($iter:expr, usize1) => { read_value!($iter, usize) - 1 }; ($iter:expr, $t:ty) => { $iter.next().unwrap().parse::<$t>().expect("Parse error") }; } // ---------- end input macro ---------- // ---------- begin recurse ---------- // reference // https://twitter.com/noshi91/status/1393952665566994434 // https://twitter.com/shino16_cp/status/1393933468082397190 pub fn recurse<A, R, F>(f: F) -> impl Fn(A) -> R where F: Fn(&dyn Fn(A) -> R, A) -> R, { fn call<A, R, F>(f: &F, a: A) -> R where F: Fn(&dyn Fn(A) -> R, A) -> R, { f(&|a| call(f, a), a) } move |a| call(&f, a) } // ---------- end recurse ---------- // ---------- begin modint ---------- use std::ops::*; pub trait Zero: Sized + Add<Self, Output = Self> { fn zero() -> Self; fn is_zero(&self) -> bool; } pub trait One: Sized + Mul<Self, Output = Self> { fn one() -> Self; fn is_one(&self) -> bool; } pub trait Ring: Zero + One + Sub<Output = Self> {} pub trait Field: Ring + Div<Output = Self> {} pub const fn pow_mod(mut r: u32, mut n: u32, m: u32) -> u32 { let mut t = 1; while n > 0 { if n & 1 == 1 { t = (t as u64 * r as u64 % m as u64) as u32; } r = (r as u64 * r as u64 % m as u64) as u32; n >>= 1; } t } pub const fn primitive_root(p: u32) -> u32 { let mut m = p - 1; let mut f = [1; 30]; let mut k = 0; let mut d = 2; while d * d <= m { if m % d == 0 { f[k] = d; k += 1; } while m % d == 0 { m /= d; } d += 1; } if m > 1 { f[k] = m; k += 1; } let mut g = 1; while g < p { let mut ok = true; let mut i = 0; while i < k { ok &= pow_mod(g, (p - 1) / f[i], p) > 1; i += 1; } if ok { break; } g += 1; } g } pub const fn is_prime(n: u32) -> bool { if n <= 1 { return false; } let mut d = 2; while d * d <= n { if n % d == 0 { return false; } d += 1; } true } #[derive(Clone, Copy, PartialEq, Eq)] pub struct ModInt<const M: u32>(u32); impl<const M: u32> ModInt<{ M }> { const REM: u32 = { let mut t = 1u32; let mut s = !M + 1; let mut n = !0u32 >> 2; while n > 0 { if n & 1 == 1 { t = t.wrapping_mul(s); } s = s.wrapping_mul(s); n >>= 1; } t }; const INI: u64 = ((1u128 << 64) % M as u128) as u64; const IS_PRIME: () = assert!(is_prime(M)); const PRIMITIVE_ROOT: u32 = primitive_root(M); const ORDER: usize = 1 << (M - 1).trailing_zeros(); const fn reduce(x: u64) -> u32 { let _ = Self::IS_PRIME; let b = (x as u32 * Self::REM) as u64; let t = x + b * M as u64; let mut c = (t >> 32) as u32; if c >= M { c -= M; } c as u32 } const fn multiply(a: u32, b: u32) -> u32 { Self::reduce(a as u64 * b as u64) } pub const fn new(v: u32) -> Self { assert!(v < M); Self(Self::reduce(v as u64 * Self::INI)) } pub const fn const_mul(&self, rhs: Self) -> Self { Self(Self::multiply(self.0, rhs.0)) } pub const fn pow(&self, mut n: u64) -> Self { let mut t = Self::new(1); let mut r = *self; while n > 0 { if n & 1 == 1 { t = t.const_mul(r); } r = r.const_mul(r); n >>= 1; } t } pub const fn inv(&self) -> Self { assert!(self.0 != 0); self.pow(M as u64 - 2) } pub const fn get(&self) -> u32 { Self::reduce(self.0 as u64) } pub const fn zero() -> Self { Self::new(0) } pub const fn one() -> Self { Self::new(1) } } impl<const M: u32> Add for ModInt<{ M }> { type Output = Self; fn add(self, rhs: Self) -> Self::Output { let mut v = self.0 + rhs.0; if v >= M { v -= M; } Self(v) } } impl<const M: u32> Sub for ModInt<{ M }> { type Output = Self; fn sub(self, rhs: Self) -> Self::Output { let mut v = self.0 - rhs.0; if self.0 < rhs.0 { v += M; } Self(v) } } impl<const M: u32> Mul for ModInt<{ M }> { type Output = Self; fn mul(self, rhs: Self) -> Self::Output { self.const_mul(rhs) } } impl<const M: u32> Div for ModInt<{ M }> { type Output = Self; fn div(self, rhs: Self) -> Self::Output { self * rhs.inv() } } impl<const M: u32> AddAssign for ModInt<{ M }> { fn add_assign(&mut self, rhs: Self) { *self = *self + rhs; } } impl<const M: u32> SubAssign for ModInt<{ M }> { fn sub_assign(&mut self, rhs: Self) { *self = *self - rhs; } } impl<const M: u32> MulAssign for ModInt<{ M }> { fn mul_assign(&mut self, rhs: Self) { *self = *self * rhs; } } impl<const M: u32> DivAssign for ModInt<{ M }> { fn div_assign(&mut self, rhs: Self) { *self = *self / rhs; } } impl<const M: u32> Neg for ModInt<{ M }> { type Output = Self; fn neg(self) -> Self::Output { if self.0 == 0 { self } else { Self(M - self.0) } } } impl<const M: u32> std::fmt::Display for ModInt<{ M }> { fn fmt<'a>(&self, f: &mut std::fmt::Formatter<'a>) -> std::fmt::Result { write!(f, "{}", self.get()) } } impl<const M: u32> std::fmt::Debug for ModInt<{ M }> { fn fmt<'a>(&self, f: &mut std::fmt::Formatter<'a>) -> std::fmt::Result { write!(f, "{}", self.get()) } } impl<const M: u32> std::str::FromStr for ModInt<{ M }> { type Err = std::num::ParseIntError; fn from_str(s: &str) -> Result<Self, Self::Err> { let val = s.parse::<u32>()?; Ok(ModInt::new(val)) } } impl<const M: u32> From<usize> for ModInt<{ M }> { fn from(val: usize) -> ModInt<{ M }> { ModInt::new((val % M as usize) as u32) } } impl<const M: u32> Zero for ModInt<{ M }> { fn zero() -> Self { Self::zero() } fn is_zero(&self) -> bool { self.0 == 0 } } impl<const M: u32> One for ModInt<{ M }> { fn one() -> Self { Self::one() } fn is_one(&self) -> bool { self.get() == 1 } } impl<const M: u32> Ring for ModInt<{ M }> {} impl<const M: u32> Field for ModInt<{ M }> {} struct NTTPrecalc<const M: u32> { sum_e: [ModInt<{ M }>; 30], sum_ie: [ModInt<{ M }>; 30], } impl<const M: u32> NTTPrecalc<{ M }> { const fn new() -> Self { let cnt2 = (M - 1).trailing_zeros() as usize; let root = ModInt::new(ModInt::<{ M }>::PRIMITIVE_ROOT); let zeta = root.pow((M - 1) as u64 >> cnt2); let mut es = [ModInt::zero(); 30]; let mut ies = [ModInt::zero(); 30]; let mut sum_e = [ModInt::zero(); 30]; let mut sum_ie = [ModInt::zero(); 30]; let mut e = zeta; let mut ie = e.inv(); let mut i = cnt2; while i >= 2 { es[i - 2] = e; ies[i - 2] = ie; e = e.const_mul(e); ie = ie.const_mul(ie); i -= 1; } let mut now = ModInt::one(); let mut inow = ModInt::one(); let mut i = 0; while i < cnt2 - 1 { sum_e[i] = es[i].const_mul(now); sum_ie[i] = ies[i].const_mul(inow); now = ies[i].const_mul(now); inow = es[i].const_mul(inow); i += 1; } Self { sum_e, sum_ie } } } struct NTTPrecalcHelper<const MOD: u32>; impl<const MOD: u32> NTTPrecalcHelper<MOD> { const A: NTTPrecalc<MOD> = NTTPrecalc::new(); } pub trait ArrayAdd { type Item; fn add(&self, rhs: &[Self::Item]) -> Vec<Self::Item>; } impl<T> ArrayAdd for [T] where T: Zero + Copy, { type Item = T; fn add(&self, rhs: &[Self::Item]) -> Vec<Self::Item> { let mut c = vec![T::zero(); self.len().max(rhs.len())]; c[..self.len()].copy_from_slice(self); c.add_assign(rhs); c } } pub trait ArrayAddAssign { type Item; fn add_assign(&mut self, rhs: &[Self::Item]); } impl<T> ArrayAddAssign for [T] where T: Add<Output = T> + Copy, { type Item = T; fn add_assign(&mut self, rhs: &[Self::Item]) { assert!(self.len() >= rhs.len()); self.iter_mut().zip(rhs).for_each(|(x, a)| *x = *x + *a); } } impl<T> ArrayAddAssign for Vec<T> where T: Zero + Add<Output = T> + Copy, { type Item = T; fn add_assign(&mut self, rhs: &[Self::Item]) { if self.len() < rhs.len() { self.resize(rhs.len(), T::zero()); } self.as_mut_slice().add_assign(rhs); } } pub trait ArraySub { type Item; fn sub(&self, rhs: &[Self::Item]) -> Vec<Self::Item>; } impl<T> ArraySub for [T] where T: Zero + Sub<Output = T> + Copy, { type Item = T; fn sub(&self, rhs: &[Self::Item]) -> Vec<Self::Item> { let mut c = vec![T::zero(); self.len().max(rhs.len())]; c[..self.len()].copy_from_slice(self); c.sub_assign(rhs); c } } pub trait ArraySubAssign { type Item; fn sub_assign(&mut self, rhs: &[Self::Item]); } impl<T> ArraySubAssign for [T] where T: Sub<Output = T> + Copy, { type Item = T; fn sub_assign(&mut self, rhs: &[Self::Item]) { assert!(self.len() >= rhs.len()); self.iter_mut().zip(rhs).for_each(|(x, a)| *x = *x - *a); } } impl<T> ArraySubAssign for Vec<T> where T: Zero + Sub<Output = T> + Copy, { type Item = T; fn sub_assign(&mut self, rhs: &[Self::Item]) { if self.len() < rhs.len() { self.resize(rhs.len(), T::zero()); } self.as_mut_slice().sub_assign(rhs); } } pub trait ArrayDot { type Item; fn dot(&self, rhs: &[Self::Item]) -> Vec<Self::Item>; } impl<T> ArrayDot for [T] where T: Mul<Output = T> + Copy, { type Item = T; fn dot(&self, rhs: &[Self::Item]) -> Vec<Self::Item> { assert!(self.len() == rhs.len()); self.iter().zip(rhs).map(|p| *p.0 * *p.1).collect() } } pub trait ArrayDotAssign { type Item; fn dot_assign(&mut self, rhs: &[Self::Item]); } impl<T> ArrayDotAssign for [T] where T: MulAssign + Copy, { type Item = T; fn dot_assign(&mut self, rhs: &[Self::Item]) { assert!(self.len() == rhs.len()); self.iter_mut().zip(rhs).for_each(|(x, a)| *x *= *a); } } pub trait ArrayMul { type Item; fn mul(&self, rhs: &[Self::Item]) -> Vec<Self::Item>; } impl<T> ArrayMul for [T] where T: Zero + One + Copy, { type Item = T; fn mul(&self, rhs: &[Self::Item]) -> Vec<Self::Item> { if self.is_empty() || rhs.is_empty() { return vec![]; } let mut res = vec![T::zero(); self.len() + rhs.len() - 1]; for (i, a) in self.iter().enumerate() { for (res, b) in res[i..].iter_mut().zip(rhs.iter()) { *res = *res + *a * *b; } } res } } // transform でlen=1を指定すればNTTになる pub trait ArrayConvolution { type Item; fn transform(&mut self, len: usize); fn inverse_transform(&mut self, len: usize); fn convolution(&self, rhs: &[Self::Item]) -> Vec<Self::Item>; } impl<const M: u32> ArrayConvolution for [ModInt<{ M }>] { type Item = ModInt<{ M }>; fn transform(&mut self, len: usize) { let f = self; let n = f.len(); let k = (n / len).trailing_zeros() as usize; assert!(len << k == n); assert!(k <= ModInt::<{ M }>::ORDER); let pre = &NTTPrecalcHelper::<M>::A; for ph in 1..=k { let p = len << (k - ph); let mut now = ModInt::one(); for (i, f) in f.chunks_exact_mut(2 * p).enumerate() { let (x, y) = f.split_at_mut(p); for (x, y) in x.iter_mut().zip(y.iter_mut()) { let l = *x; let r = *y * now; *x = l + r; *y = l - r; } now *= pre.sum_e[(!i).trailing_zeros() as usize]; } } } fn inverse_transform(&mut self, len: usize) { let f = self; let n = f.len(); let k = (n / len).trailing_zeros() as usize; assert!(len << k == n); assert!(k <= ModInt::<{ M }>::ORDER); let pre = &NTTPrecalcHelper::<M>::A; for ph in (1..=k).rev() { let p = len << (k - ph); let mut inow = ModInt::one(); for (i, f) in f.chunks_exact_mut(2 * p).enumerate() { let (x, y) = f.split_at_mut(p); for (x, y) in x.iter_mut().zip(y.iter_mut()) { let l = *x; let r = *y; *x = l + r; *y = (l - r) * inow; } inow *= pre.sum_ie[(!i).trailing_zeros() as usize]; } } let ik = ModInt::new(2).inv().pow(k as u64); for f in f.iter_mut() { *f *= ik; } } fn convolution(&self, rhs: &[Self::Item]) -> Vec<Self::Item> { if self.len().min(rhs.len()) <= 32 { return self.mul(rhs); } const PARAM: usize = 10; let size = self.len() + rhs.len() - 1; let mut k = 0; while (size + (1 << k) - 1) >> k > PARAM { k += 1; } let len = (size + (1 << k) - 1) >> k; let mut f = vec![ModInt::zero(); len << k]; let mut g = vec![ModInt::zero(); len << k]; f[..self.len()].copy_from_slice(self); g[..rhs.len()].copy_from_slice(rhs); f.transform(len); g.transform(len); let mut buf = [ModInt::zero(); 2 * PARAM - 1]; let buf = &mut buf[..(2 * len - 1)]; let pre = &NTTPrecalcHelper::<M>::A; let mut now = ModInt::one(); for (i, (f, g)) in f .chunks_exact_mut(2 * len) .zip(g.chunks_exact(2 * len)) .enumerate() { let mut r = now; for (f, g) in f.chunks_exact_mut(len).zip(g.chunks_exact(len)) { buf.fill(ModInt::zero()); for (i, f) in f.iter().enumerate() { for (buf, g) in buf[i..].iter_mut().zip(g.iter()) { *buf = *buf + *f * *g; } } f.copy_from_slice(&buf[..len]); for (f, buf) in f.iter_mut().zip(buf[len..].iter()) { *f = *f + r * *buf; } r = -r; } now *= pre.sum_e[(!i).trailing_zeros() as usize]; } f.inverse_transform(len); f.truncate(self.len() + rhs.len() - 1); f } } // ---------- end modint ----------