結果
問題 | No.2116 Making Forest Hard |
ユーザー | akakimidori |
提出日時 | 2024-01-29 23:20:27 |
言語 | Rust (1.77.0 + proconio) |
結果 |
AC
|
実行時間 | 1,833 ms / 8,000 ms |
コード長 | 27,285 bytes |
コンパイル時間 | 14,430 ms |
コンパイル使用メモリ | 378,032 KB |
実行使用メモリ | 40,576 KB |
最終ジャッジ日時 | 2024-09-28 10:07:39 |
合計ジャッジ時間 | 51,792 ms |
ジャッジサーバーID (参考情報) |
judge3 / judge1 |
(要ログイン)
テストケース
テストケース表示入力 | 結果 | 実行時間 実行使用メモリ |
---|---|---|
testcase_00 | AC | 1 ms
5,248 KB |
testcase_01 | AC | 1 ms
5,376 KB |
testcase_02 | AC | 1,833 ms
40,536 KB |
testcase_03 | AC | 1,821 ms
40,356 KB |
testcase_04 | AC | 676 ms
30,712 KB |
testcase_05 | AC | 640 ms
31,148 KB |
testcase_06 | AC | 741 ms
30,648 KB |
testcase_07 | AC | 756 ms
30,812 KB |
testcase_08 | AC | 556 ms
23,208 KB |
testcase_09 | AC | 605 ms
26,760 KB |
testcase_10 | AC | 44 ms
5,376 KB |
testcase_11 | AC | 459 ms
26,072 KB |
testcase_12 | AC | 570 ms
28,504 KB |
testcase_13 | AC | 586 ms
33,008 KB |
testcase_14 | AC | 1,177 ms
31,820 KB |
testcase_15 | AC | 10 ms
5,376 KB |
testcase_16 | AC | 676 ms
30,204 KB |
testcase_17 | AC | 811 ms
24,688 KB |
testcase_18 | AC | 1,441 ms
35,416 KB |
testcase_19 | AC | 283 ms
13,936 KB |
testcase_20 | AC | 320 ms
15,220 KB |
testcase_21 | AC | 450 ms
15,852 KB |
testcase_22 | AC | 516 ms
28,900 KB |
testcase_23 | AC | 75 ms
6,016 KB |
testcase_24 | AC | 177 ms
6,492 KB |
testcase_25 | AC | 428 ms
18,048 KB |
testcase_26 | AC | 736 ms
30,224 KB |
testcase_27 | AC | 149 ms
7,808 KB |
testcase_28 | AC | 1,145 ms
27,216 KB |
testcase_29 | AC | 522 ms
29,164 KB |
testcase_30 | AC | 1,425 ms
33,400 KB |
testcase_31 | AC | 156 ms
11,136 KB |
testcase_32 | AC | 646 ms
33,044 KB |
testcase_33 | AC | 25 ms
5,376 KB |
testcase_34 | AC | 191 ms
9,600 KB |
testcase_35 | AC | 684 ms
33,208 KB |
testcase_36 | AC | 535 ms
23,660 KB |
testcase_37 | AC | 29 ms
5,376 KB |
testcase_38 | AC | 679 ms
30,284 KB |
testcase_39 | AC | 639 ms
30,760 KB |
testcase_40 | AC | 142 ms
10,208 KB |
testcase_41 | AC | 467 ms
25,244 KB |
testcase_42 | AC | 1,471 ms
35,756 KB |
testcase_43 | AC | 1,821 ms
40,540 KB |
testcase_44 | AC | 601 ms
29,508 KB |
testcase_45 | AC | 278 ms
17,684 KB |
testcase_46 | AC | 1,778 ms
40,576 KB |
testcase_47 | AC | 157 ms
9,984 KB |
testcase_48 | AC | 151 ms
9,316 KB |
testcase_49 | AC | 1 ms
5,376 KB |
testcase_50 | AC | 613 ms
33,044 KB |
testcase_51 | AC | 613 ms
31,200 KB |
testcase_52 | AC | 615 ms
31,052 KB |
testcase_53 | AC | 586 ms
30,648 KB |
testcase_54 | AC | 1,143 ms
31,876 KB |
ソースコード
// dp[v][a]: 部分木vでvを含みそれが最上位な連結成分について // 連結成分のmaxがa であるようなものについての // ... 何を管理するとdpできる? // way: そうなるような切り方の個数 // sum: そうなる切り方についてのvを含む連結成分サイズの和 // というのを管理できてるとして // 根vに親pが生成される時 // dp[p][A_p].0 = dp[v][*].0 + dp[v][A <= A_p].0 // dp[p][A_p].1 = 2 * dp[v][*].0 + dp[v][A <= A_p].1 // // dp[p][x] = (0, 0) (x < A_p) // dp[p][x].0 = dp[v][x].0 // dp[p][x].1 = dp[v][x].0 + dp[v][x].1 (x > A_p) // // 子にuが生えた時 type M = ModInt<998244353>; fn run() { input! { n: usize, a: [u32; n], e: [(usize1, usize1); n - 1], } let a = a .into_iter() .enumerate() .map(|p| (p.1, p.0)) .collect::<Vec<_>>(); let mut g = vec![vec![]; n]; for &(a, b) in e.iter() { g[a].push(b); g[b].push(a); } let root = 0; 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); topo.push(u); parent[u] = v; } } let mut size = vec![1usize; n]; for &v in topo.iter().rev() { g[v].sort_by_key(|u| !size[*u]); for &u in g[v].iter() { size[v] += size[u]; } } let mut memo: Vec<Vec<((u32, usize), Dual<M>)>> = vec![vec![]; n]; let mut ans = M::zero(); for &v in topo.iter().rev() { if v != root && g[parent[v]][0] == v { continue; } let mut z = vec![]; let mut pos = v; let mut path = vec![]; loop { z.push(a[pos]); path.push(pos); for &u in g[pos].iter().skip(1) { z.extend(memo[u].iter().map(|p| p.0)); } if g[pos].is_empty() { break; } pos = g[pos][0]; } z.sort(); let mut seg = LazySegmentTree::new(z.len(), R); for &v in path.iter().rev() { let x = z.binary_search(&a[v]).unwrap(); let Dual(w, s) = seg.find(0, x).0; let Dual(all, _) = seg.find(0, z.len()).0; seg.update(0, x, Dual::zero()); seg.update(x + 1, z.len(), Dual::new(M::one(), M::one())); let mut p = Dual::new(all + w, all + s + w); if v == *path.last().unwrap() { p = Dual::new(M::one(), M::one()); } seg.set_at(x, (p, p * Dual::new(M::new(a[v].0), M::zero()))); for &u in g[v].iter().skip(1) { let c = std::mem::take(&mut memo[u]); let mut memo = vec![Dual::zero(); c.len()]; let mut pos = vec![0; c.len()]; for (i, &(key, _)) in c.iter().enumerate() { let x = z.binary_search(&key).unwrap(); memo[i] = seg.find(0, x).0; pos[i] = x; } pos.insert(0, 0); pos.push(z.len()); let way = c.iter().fold(M::zero(), |s, &(_, a)| s + a.0); let mut pre = Dual::new(way, M::zero()); for (i, x) in pos.windows(2).enumerate() { let (l, r) = (x[0], x[1]); seg.update(l, r, pre); if i < c.len() { let c = c[i].1; pre = pre + c; } } for i in 0..c.len() { let a = M::new(c[i].0 .0); let c = c[i].1; let memo = memo[i]; let pos = pos[i + 1]; seg.set_at(pos, (c * memo, c * memo * Dual::new(a, M::zero()))); } } if v == root { ans += seg.find(0, z.len()).1 .1; } else { ans += seg.find(0, z.len()).1 .1 * M::new(2).pow((n - 1 - size[v]) as u64); } } memo[v] = (0..z.len()).map(|x| (z[x], seg.find(x, x + 1).0)).collect(); } println!("{}", ans); } struct R; impl TE for R { type T = (Dual<M>, Dual<M>); type E = Dual<M>; fn fold(&self, l: &Self::T, r: &Self::T) -> Self::T { (l.0 + r.0, l.1 + r.1) } fn eval(&self, x: &Self::T, f: &Self::E) -> Self::T { (x.0 * *f, x.1 * *f) } fn merge(&self, g: &Self::E, h: &Self::E) -> Self::E { *g * *h } fn e(&self) -> Self::T { (Dual::zero(), Dual::zero()) } fn id(&self) -> Self::E { Dual::one() } } fn main() { run(); } #[derive(Clone, Copy, Default, Debug)] pub struct Dual<T>(T, T); impl<T> Dual<T> { pub fn new(a: T, b: T) -> Self { Self(a, b) } } impl<T> Zero for Dual<T> where T: Zero, { fn zero() -> Self { Self::new(T::zero(), T::zero()) } fn is_zero(&self) -> bool { self.0.is_zero() && self.1.is_zero() } } impl<T> One for Dual<T> where T: One + Zero + Clone, { fn one() -> Self { Self::new(T::one(), T::zero()) } fn is_one(&self) -> bool { self.0.is_one() && self.1.is_zero() } } impl<T> Add for Dual<T> where T: Add<Output = T>, { type Output = Self; fn add(self, rhs: Self) -> Self { Self::new(self.0 + rhs.0, self.1 + rhs.1) } } impl<T> Sub for Dual<T> where T: Sub<Output = T>, { type Output = Self; fn sub(self, rhs: Self) -> Self { Self::new(self.0 - rhs.0, self.1 - rhs.1) } } impl<T> Mul for Dual<T> where T: Clone + Add<Output = T> + Mul<Output = T>, { type Output = Self; fn mul(self, rhs: Self) -> Self { Self::new( self.0.clone() * rhs.0.clone(), self.0 * rhs.1 + self.1 * rhs.0, ) } } // ---------- 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 Lazy Segment Tree ---------- pub trait TE { type T: Clone; type E: Clone; fn fold(&self, l: &Self::T, r: &Self::T) -> Self::T; fn eval(&self, x: &Self::T, f: &Self::E) -> Self::T; fn merge(&self, g: &Self::E, h: &Self::E) -> Self::E; fn e(&self) -> Self::T; fn id(&self) -> Self::E; } pub struct LazySegmentTree<R: TE> { n: usize, size: usize, bit: u32, op: R, data: Vec<(R::T, R::E)>, } impl<R: TE> LazySegmentTree<R> { pub fn new(n: usize, op: R) -> Self { assert!(n > 0); let size = n.next_power_of_two(); let bit = size.trailing_zeros(); let data = vec![(op.e(), op.id()); 2 * size]; Self { n, size, bit, op, data, } } pub fn build<I>(init: I, n: usize, op: R) -> Self where I: Iterator<Item = R::T>, { let mut seg = Self::new(n, op); for (data, ini) in seg.data[seg.size..].iter_mut().zip(init) { data.0 = ini; } for i in (1..seg.size).rev() { seg.pull(i); } seg } pub fn update(&mut self, l: usize, r: usize, f: R::E) { assert!(l <= r && r <= self.n); if l == r { return; } self.push_range(l, r); let mut s = l + self.size; let mut t = r + self.size; while s < t { if s & 1 == 1 { self.apply(s, &f); s += 1; } if t & 1 == 1 { t -= 1; self.apply(t, &f); } s >>= 1; t >>= 1; } let l = l + self.size; let r = r + self.size; for k in 1..=self.bit { if (l >> k) << k != l { self.pull(l >> k); } if (r >> k) << k != r { self.pull((r - 1) >> k); } } } pub fn find(&mut self, l: usize, r: usize) -> R::T { assert!(l <= r && r <= self.n); if l == r { return self.op.e(); } self.push_range(l, r); let mut l = l + self.size; let mut r = r + self.size; let mut p = self.op.e(); let mut q = self.op.e(); while l < r { if l & 1 == 1 { p = self.op.fold(&p, &self.data[l].0); l += 1; } if r & 1 == 1 { r -= 1; q = self.op.fold(&self.data[r].0, &q); } l >>= 1; r >>= 1; } self.op.fold(&p, &q) } pub fn set_at(&mut self, x: usize, v: R::T) { assert!(x < self.n); let x = x + self.size; for k in (1..=self.bit).rev() { self.push(x >> k); } self.data[x].0 = v; for k in 1..=self.bit { self.pull(x >> k); } } fn push_range(&mut self, l: usize, r: usize) { let l = l + self.size; let r = r + self.size; for k in (1..(self.bit + 1)).rev() { if (l >> k) << k != l { self.push(l >> k); } if (r >> k) << k != r { self.push((r - 1) >> k); } } } fn apply(&mut self, x: usize, f: &R::E) { self.data[x].0 = self.op.eval(&self.data[x].0, f); self.data[x].1 = self.op.merge(&self.data[x].1, f); } fn push(&mut self, x: usize) { let f = std::mem::replace(&mut self.data[x].1, self.op.id()); self.apply(2 * x, &f); self.apply(2 * x + 1, &f); } fn pull(&mut self, x: usize) { self.data[x].0 = self.op.fold(&self.data[2 * x].0, &self.data[2 * x + 1].0); } } // ---------- end Lazy Segment Tree ---------- use std::ops::*; // ---------- begin trait ---------- 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 SemiRing: Zero + One {} pub trait Ring: SemiRing + Sub<Output = Self> + Neg<Output = Self> {} pub trait Field: Ring + Div<Output = Self> {} impl<T> SemiRing for T where T: Zero + One {} impl<T> Ring for T where T: SemiRing + Sub<Output = Self> + Neg<Output = Self> {} impl<T> Field for T where T: Ring + Div<Output = Self> {} // ---------- end trait ---------- // ---------- begin modint ---------- 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) } } // ---------- end modint ---------- // ---------- begin precalc ---------- pub struct Precalc<const MOD: u32> { fact: Vec<ModInt<MOD>>, ifact: Vec<ModInt<MOD>>, inv: Vec<ModInt<MOD>>, } impl<const MOD: u32> Precalc<MOD> { pub fn new(size: usize) -> Self { let mut fact = vec![ModInt::one(); size + 1]; let mut ifact = vec![ModInt::one(); size + 1]; let mut inv = vec![ModInt::one(); size + 1]; for i in 2..=size { fact[i] = fact[i - 1] * ModInt::from(i); } ifact[size] = fact[size].inv(); for i in (2..=size).rev() { inv[i] = ifact[i] * fact[i - 1]; ifact[i - 1] = ifact[i] * ModInt::from(i); } Self { fact, ifact, inv } } pub fn fact(&self, n: usize) -> ModInt<MOD> { self.fact[n] } pub fn ifact(&self, n: usize) -> ModInt<MOD> { self.ifact[n] } pub fn inv(&self, n: usize) -> ModInt<MOD> { assert!(0 < n); self.inv[n] } pub fn perm(&self, n: usize, k: usize) -> ModInt<MOD> { if k > n { return ModInt::zero(); } self.fact[n] * self.ifact[n - k] } pub fn binom(&self, n: usize, k: usize) -> ModInt<MOD> { if n < k { return ModInt::zero(); } self.fact[n] * self.ifact[k] * self.ifact[n - k] } } // ---------- end precalc ---------- 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 } } // ---------- begin array op ---------- 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 array op ----------