結果
問題 | No.235 めぐるはめぐる (5) |
ユーザー | ngtkana |
提出日時 | 2024-06-24 13:40:09 |
言語 | Rust (1.77.0 + proconio) |
結果 |
AC
|
実行時間 | 1,778 ms / 10,000 ms |
コード長 | 55,598 bytes |
コンパイル時間 | 14,185 ms |
コンパイル使用メモリ | 380,288 KB |
実行使用メモリ | 50,876 KB |
最終ジャッジ日時 | 2024-06-24 13:40:33 |
合計ジャッジ時間 | 21,830 ms |
ジャッジサーバーID (参考情報) |
judge5 / judge3 |
(要ログイン)
テストケース
テストケース表示入力 | 結果 | 実行時間 実行使用メモリ |
---|---|---|
testcase_00 | AC | 1,778 ms
49,792 KB |
testcase_01 | AC | 1,148 ms
50,876 KB |
testcase_02 | AC | 1,681 ms
49,920 KB |
コンパイルメッセージ
warning: unused import: `factorial::Factorial` --> src/main.rs:1102:13 | 1102 | pub use factorial::Factorial; | ^^^^^^^^^^^^^^^^^^^^ | = note: `#[warn(unused_imports)]` on by default warning: unused import: `fourier::any_mod_fps_mul` --> src/main.rs:1103:13 | 1103 | pub use fourier::any_mod_fps_mul; | ^^^^^^^^^^^^^^^^^^^^^^^^ warning: unused import: `fourier::fft` --> src/main.rs:1104:13 | 1104 | pub use fourier::fft; | ^^^^^^^^^^^^ warning: unused import: `fourier::fps_mul` --> src/main.rs:1105:13 | 1105 | pub use fourier::fps_mul; | ^^^^^^^^^^^^^^^^ warning: unused import: `fourier::ifft` --> src/main.rs:1106:13 | 1106 | pub use fourier::ifft; | ^^^^^^^^^^^^^
ソースコード
use lazy_segtree::LazySegtree; use proconio::input; use proconio::marker::Usize1; type Fp = fp::Fp<1_000_000_007>; fn main() { input! { n: usize, sum: [usize; n], coeff: [usize; n], edges: [(Usize1, Usize1); n - 1], q: usize, } let (hld, _g) = Hld::from_edges(0, &edges); let mut values = vec![ Value { coeff: fp!(0), sum: fp!(0) }; n ]; for (x, &i) in hld.index.iter().enumerate() { values[i] = Value { coeff: fp!(coeff[x]), sum: fp!(sum[x]), }; } let mut lazy_segtree = LazySegtree::<O>::from_iter(values); for _ in 0..q { input! { com: String, } match com.as_str() { "0" => { input! { i: Usize1, j: Usize1, op: usize, } for (i, j) in hld.iter_path_segments_including_lca_by_index(i, j) { lazy_segtree.range_apply(i..=j, &fp!(op)); } } "1" => { input! { i: Usize1, j: Usize1, } let ans = hld .iter_path_segments_including_lca_by_index(i, j) .map(|(i, j)| lazy_segtree.fold(i..=j).sum) .sum::<Fp>(); println!("{ans}"); } _ => unreachable!(), } } } #[derive(Clone, Copy, Debug, PartialEq)] struct Value { coeff: Fp, sum: Fp, } enum O {} impl lazy_segtree::Op for O { type Operator = Fp; type Value = Value; fn identity() -> Self::Value { Value { coeff: fp!(0), sum: fp!(0), } } fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value { Value { coeff: lhs.coeff + rhs.coeff, sum: lhs.sum + rhs.sum, } } fn apply(op: &Self::Operator, value: &Self::Value) -> Self::Value { Value { coeff: value.coeff, sum: value.sum + op * value.coeff, } } fn identity_op() -> Self::Operator { fp!(0) } fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator { op + other } } pub struct Hld { pub parent: Vec<usize>, pub index: Vec<usize>, pub head: Vec<usize>, } impl Hld { pub fn from_short_parents(mut parent: Vec<usize>) -> (Self, Vec<Vec<usize>>) { parent.insert(0, 0); let mut g = vec![Vec::new(); parent.len()]; for (i, &p) in parent.iter().enumerate().skip(1) { g[p].push(i); } (__build_hld(0, &mut g, parent), g) } pub fn from_edges(root: usize, edges: &[(usize, usize)]) -> (Self, Vec<Vec<usize>>) { let mut g = vec![Vec::new(); edges.len() + 1]; for &(i, j) in edges { g[i].push(j); g[j].push(i); } let parent = __remove_parent(root, &mut g); (__build_hld(root, &mut g, parent), g) } pub fn visit_path_segments_including_lca( &self, mut i: usize, // id mut j: usize, // id mut f: impl FnMut(usize, usize), // id ) { while self.head[i] != self.head[j] { if self.index[i] < self.index[j] { f(self.head[j], j); j = self.parent[self.head[j]]; } else { f(self.head[i], i); i = self.parent[self.head[i]]; } } if self.index[i] < self.index[j] { f(i, j) } else { f(j, i) } } pub fn iter_path_segments_including_lca(&self, i: usize, j: usize) -> IterV<'_> { IterV { hld: self, i, j, exhausted: false, } } pub fn visit_path_segments_including_lca_by_index( &self, i: usize, // id j: usize, // id mut f: impl FnMut(usize, usize), // index ) { self.visit_path_segments_including_lca(i, j, |i, j| f(self.index[i], self.index[j])); } pub fn iter_path_segments_including_lca_by_index( &self, i: usize, j: usize, ) -> impl Iterator<Item = (usize, usize)> + '_ { self.iter_path_segments_including_lca(i, j) .map(|(i, j)| (self.index[i], self.index[j])) } pub fn ledacy_iter_v(&self, i: usize, j: usize) -> impl Iterator<Item = (usize, usize)> + '_ { self.iter_path_segments_including_lca_by_index(i, j) } pub fn lca(&self, mut i: usize, mut j: usize) -> usize { while self.head[i] != self.head[j] { if self.index[i] < self.index[j] { j = self.parent[self.head[j]]; } else { i = self.parent[self.head[i]]; } } std::cmp::min_by_key(i, j, |&i| self.index[i]) } } pub struct IterV<'a> { hld: &'a Hld, i: usize, j: usize, exhausted: bool, } impl Iterator for IterV<'_> { type Item = (usize, usize); fn next(&mut self) -> Option<Self::Item> { (!self.exhausted).then_some(())?; let Self { hld, i, j, .. } = *self; let Hld { index, head, parent, } = hld; Some(if head[i] == head[j] { self.exhausted = true; if index[i] < index[j] { (i, j) } else { (j, i) } } else { if index[i] < index[j] { self.j = parent[head[j]]; (head[j], j) } else { self.i = parent[head[i]]; (head[i], i) } }) } } fn __build_hld(root: usize, g: &mut [Vec<usize>], parent: Vec<usize>) -> Hld { let n = g.len(); __heavy_first(0, g); let mut index = vec![usize::MAX; n]; let mut head = vec![usize::MAX; n]; head[root] = root; __head_and_index(0, &*g, &mut head, &mut index, &mut (0..)); Hld { parent, index, head, } } fn __head_and_index( i: usize, g: &[Vec<usize>], head: &mut [usize], index: &mut [usize], current: &mut std::ops::RangeFrom<usize>, ) { index[i] = current.next().unwrap(); for &j in &g[i] { head[j] = if j == g[i][0] { head[i] } else { j }; __head_and_index(j, g, head, index, current); } } fn __heavy_first(i: usize, g: &mut [Vec<usize>]) -> usize { let mut max = 0; let mut size = 1; for e in 0..g[i].len() { let csize = __heavy_first(g[i][e], g); if max < csize { max = csize; g[i].swap(0, e); } size += csize; } size } fn __remove_parent(root: usize, g: &mut [Vec<usize>]) -> Vec<usize> { let mut stack = vec![root]; let mut parent = vec![usize::MAX; g.len()]; parent[root] = root; while let Some(i) = stack.pop() { g[i].retain(|&j| parent[i] != j); for &j in &g[i] { parent[j] = i; stack.push(j); } } parent } // link_cut_tree {{{ // https://ngtkana.github.io/ac-adapter-rs/link_cut_tree/index.html #[allow(dead_code)] mod link_cut_tree { mod base { #[doc(hidden)] pub trait OpBase { type Value: Clone; type InternalValue: Clone; fn identity() -> Self::InternalValue; fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue; fn into_front(value: Self::InternalValue) -> Self::Value; fn from_front(value: Self::Value) -> Self::InternalValue; fn rev(value: &mut Self::InternalValue); } pub struct LinkCutTreeBase<O: OpBase> { nodes: Vec<Node<O>>, } impl<O: OpBase> LinkCutTreeBase<O> { pub fn new(n: usize) -> Self { Self { nodes: (0..n) .map(|id| Node { id, parent: std::ptr::null_mut(), left: std::ptr::null_mut(), right: std::ptr::null_mut(), rev: false, value: O::identity(), acc: O::identity(), }) .collect(), } } pub fn from_values(values: impl IntoIterator<Item = O::Value>) -> Self { Self { nodes: values .into_iter() .map(O::from_front) .enumerate() .map(|(id, value)| Node { id, parent: std::ptr::null_mut(), left: std::ptr::null_mut(), right: std::ptr::null_mut(), rev: false, value: value.clone(), acc: value, }) .collect(), } } pub fn link(&mut self, p: usize, c: usize) { unsafe { let c = std::ptr::addr_of_mut!(self.nodes[c]); let p = std::ptr::addr_of_mut!(self.nodes[p]); expose(c); assert!((*c).left.is_null(), "c = {} is not a root", (*c).id); expose(p); assert!( (*c).parent.is_null(), "c = {} and p = {} are already connected", (*c).id, (*p).id ); (*c).parent = p; (*p).right = c; update(p); } } pub fn undirected_link(&mut self, i: usize, j: usize) -> bool { if self.undirected_is_connected(i, j) { return false; } self.evert(j); self.link(i, j); true } pub fn cut(&mut self, x: usize) -> Option<usize> { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); let p = (*x).left; (*x).left = std::ptr::null_mut(); let ans = p.as_ref().map(|p| p.id); if !p.is_null() { (*p).parent = std::ptr::null_mut(); } update(x); ans } } pub fn undirected_cut(&mut self, i: usize, j: usize) -> bool { if !self.undirected_has_edge(i, j) { return false; } self.evert(i); self.cut(j); true } pub fn evert(&mut self, x: usize) { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); rev(x); push(x); } } pub fn undirected_has_edge(&mut self, x: usize, y: usize) -> bool { self.parent(x) == Some(y) || self.parent(y) == Some(x) } pub fn undirected_is_connected(&mut self, x: usize, y: usize) -> bool { if x == y { return true; } unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); let y = std::ptr::addr_of_mut!(self.nodes[y]); expose(x); expose(y); !(*x).parent.is_null() } } pub fn lca(&mut self, x: usize, y: usize) -> Option<usize> { if x == y { return Some(x); } unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); let y = std::ptr::addr_of_mut!(self.nodes[y]); expose(x); let lca = expose(y); if (*x).parent.is_null() { None } else { Some((*lca).id) } } } pub fn set(&mut self, x: usize, mut f: impl FnMut(O::Value) -> O::Value) { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); (*x).value = O::from_front(f(O::into_front((*x).value.clone()))); update(x); } } pub fn fold(&mut self, x: usize) -> O::Value { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); O::into_front((*x).acc.clone()) } } pub fn undirected_fold(&mut self, i: usize, j: usize) -> Option<O::Value> { if !self.undirected_is_connected(i, j) { return None; } self.evert(i); Some(self.fold(j)) } pub fn parent(&mut self, x: usize) -> Option<usize> { unsafe { let x = std::ptr::addr_of_mut!(self.nodes[x]); expose(x); let mut p = (*x).left.as_mut()?; while let Some(next) = p.right.as_mut() { p = next; } splay(p); Some(p.id) } } } #[derive(Clone, Copy)] struct Node<O: OpBase> { id: usize, parent: *mut Self, left: *mut Self, right: *mut Self, rev: bool, value: O::InternalValue, acc: O::InternalValue, } unsafe fn is_splay_root<O: OpBase>(x: *mut Node<O>) -> bool { let x = &*x; let p = match x.parent.as_ref() { Some(p) => p, None => return true, }; !std::ptr::eq(x, p.left) && !std::ptr::eq(x, p.right) } unsafe fn push<O: OpBase>(x: *mut Node<O>) { let x = &mut *x; if x.rev { if let Some(l) = x.left.as_mut() { rev(l); } if let Some(r) = x.right.as_mut() { rev(r); } x.rev = false; } } unsafe fn update<O: OpBase>(x: *mut Node<O>) { let x = &mut *x; x.acc = x.value.clone(); if !x.left.is_null() { x.acc = O::mul(&(*x.left).acc, &x.acc); } if !x.right.is_null() { x.acc = O::mul(&x.acc, &(*x.right).acc); } } unsafe fn rev<O: OpBase>(x: *mut Node<O>) { let x = &mut *x; std::mem::swap(&mut x.left, &mut x.right); O::rev(&mut x.acc); x.rev ^= true; } unsafe fn expose<O: OpBase>(x: *mut Node<O>) -> *mut Node<O> { let mut last = std::ptr::null_mut(); let mut current = x; while !current.is_null() { splay(current); (*current).right = last; update(current); last = current; current = (*current).parent; } splay(x); last } unsafe fn splay<O: OpBase>(x: *mut Node<O>) { let x = &mut *x; push(x); while !is_splay_root(x) { let p = &mut *x.parent; if is_splay_root(p) { push(p); push(x); if std::ptr::eq(p.left, x) { rotate_right(p); } else { rotate_left(p); } } else { let g = &mut *p.parent; push(g); push(p); push(x); #[allow(clippy::collapsible_else_if)] if std::ptr::eq(p.left, x) { if std::ptr::eq(g.left, p) { rotate_right(g); rotate_right(p); } else { rotate_right(p); rotate_left(g); } } else { if std::ptr::eq(g.left, p) { rotate_left(p); rotate_right(g); } else { rotate_left(g); rotate_left(p); } } } } } unsafe fn rotate_left<O: OpBase>(l: *mut Node<O>) { let l = &mut *l; let r = &mut *l.right; let p = l.parent; let c = r.left; l.right = c; if !c.is_null() { (*c).parent = l; } r.left = l; l.parent = r; r.parent = p; update(l); update(r); if !p.is_null() { if std::ptr::eq((*p).left, l) { (*p).left = r; } else if std::ptr::eq((*p).right, l) { (*p).right = r; } update(&mut *p); } } unsafe fn rotate_right<O: OpBase>(r: *mut Node<O>) { let r = &mut *r; let l = &mut *r.left; let p = r.parent; let c = l.right; r.left = c; if !c.is_null() { (*c).parent = r; } l.right = r; r.parent = l; l.parent = p; update(r); update(l); if !p.is_null() { if std::ptr::eq((*p).left, r) { (*p).left = l; } else if std::ptr::eq((*p).right, r) { (*p).right = l; } update(&mut *p); } } } pub use base::LinkCutTreeBase; use base::OpBase; pub trait Op { type Value: Clone; fn identity() -> Self::Value; fn mul(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value; } impl OpBase for () { type InternalValue = (); type Value = (); fn identity() -> Self::InternalValue {} fn mul(_lhs: &Self::InternalValue, _rhs: &Self::InternalValue) -> Self::InternalValue {} fn rev(_value: &mut Self::InternalValue) {} fn into_front(_value: Self::InternalValue) {} fn from_front(_value: Self::Value) -> Self::InternalValue {} } pub type LinkCutTree = LinkCutTreeBase<()>; pub type CommutLinkCutTree<T> = LinkCutTreeBase<Commut<T>>; #[doc(hidden)] pub struct Commut<T: Op>(T); impl<T: Op> OpBase for Commut<T> { type InternalValue = T::Value; type Value = T::Value; fn identity() -> Self::InternalValue { T::identity() } fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue { T::mul(lhs, rhs) } fn rev(_value: &mut Self::InternalValue) {} fn into_front(value: Self::InternalValue) -> Self::Value { value } fn from_front(value: Self::Value) -> Self::InternalValue { value } } #[doc(hidden)] pub struct NonCommut<T: Op>(T); pub type NonCommutLinkCutTree<T> = LinkCutTreeBase<NonCommut<T>>; impl<T: Op> OpBase for NonCommut<T> { type InternalValue = (T::Value, T::Value); type Value = T::Value; fn identity() -> Self::InternalValue { (T::identity(), T::identity()) } fn mul(lhs: &Self::InternalValue, rhs: &Self::InternalValue) -> Self::InternalValue { (T::mul(&lhs.0, &rhs.0), T::mul(&rhs.1, &lhs.1)) } fn rev(value: &mut Self::InternalValue) { std::mem::swap(&mut value.0, &mut value.1); } fn into_front(value: Self::InternalValue) -> Self::Value { value.0 } fn from_front(value: Self::Value) -> Self::InternalValue { (value.clone(), value) } } } // }}} // lazy_segtree {{{ // https://ngtkana.github.io/ac-adapter-rs/lazy_segtree/index.html #[allow(dead_code)] mod lazy_segtree { use std::iter::FromIterator; use std::mem::replace; use std::ops::RangeBounds; pub trait Op { type Value; type Operator: PartialEq; fn identity() -> Self::Value; fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value; fn apply(op: &Self::Operator, value: &Self::Value) -> Self::Value; fn identity_op() -> Self::Operator; fn compose(op: &Self::Operator, other: &Self::Operator) -> Self::Operator; } #[derive(Debug, Clone)] pub struct LazySegtree<O: Op> { values: Vec<O::Value>, operators: Vec<O::Operator>, } impl<O: Op> LazySegtree<O> { pub fn new(values: &[O::Value]) -> Self where O::Value: Clone, O::Operator: Clone, { let values_ = values; let n = values_.len(); let mut values = vec![O::identity(); 2 * n]; values[n..].clone_from_slice(values_); for i in (1..n).rev() { values[i] = O::op(&values[i * 2], &values[i * 2 + 1]); } Self { values, operators: vec![O::identity_op(); 2 * n], } } pub fn range_apply<R: RangeBounds<usize>>(&mut self, range: R, f: &O::Operator) { let n = self.operators.len() / 2; let (l, r) = open(range, n); if l == r { return; } let l = n + l; let r = n + r; for p in (0..usize::BITS - r.leading_zeros()).rev() { self.push(l >> p); self.push((r - 1) >> p); } { let mut l = l; let mut r = r; while l < r { if l & 1 != 0 { self.operators[l] = O::compose(f, &self.operators[l]); l += 1; } if r & 1 != 0 { r -= 1; self.operators[r] = O::compose(f, &self.operators[r]); } l >>= 1; r >>= 1; } } for p in 1..usize::BITS - r.leading_zeros() { self.update(l >> p); self.update((r - 1) >> p); } } pub fn fold<R: RangeBounds<usize>>(&mut self, range: R) -> O::Value { let n = self.operators.len() / 2; let (mut l, mut r) = open(range, n); if l == r { return O::identity(); } l += n; r += n; for p in (0..usize::BITS - r.leading_zeros()).rev() { self.push(l >> p); self.push((r - 1) >> p); } for p in 1..usize::BITS - r.leading_zeros() { self.update(l >> p); self.update((r - 1) >> p); } let mut left = O::identity(); let mut right = O::identity(); while l < r { if l & 1 != 0 { left = O::op(&left, &O::apply(&self.operators[l], &self.values[l])); l += 1; } if r & 1 != 0 { r -= 1; right = O::op(&O::apply(&self.operators[r], &self.values[r]), &right); } l >>= 1; r >>= 1; } O::op(&left, &right) } pub fn get(&self, i: usize) -> O::Value where O::Value: Clone, { let mut i = self.operators.len() / 2 + i; let mut value = self.values[i].clone(); while i > 0 { value = O::apply(&self.operators[i], &value); i >>= 1; } value } pub fn collect(&self) -> Vec<O::Value> where O::Value: Clone, { (0..self.operators.len() / 2).map(|i| self.get(i)).collect() } fn push(&mut self, i: usize) { let a = replace(&mut self.operators[i], O::identity_op()); self.values[i] = O::apply(&a, &self.values[i]); if i < self.operators.len() / 2 { self.operators[i << 1] = O::compose(&a, &self.operators[i << 1]); self.operators[i << 1 | 1] = O::compose(&a, &self.operators[i << 1 | 1]); } } fn update(&mut self, i: usize) { self.values[i] = O::op( &O::apply(&self.operators[i << 1], &self.values[i << 1]), &O::apply(&self.operators[i << 1 | 1], &self.values[i << 1 | 1]), ) } } impl<O: Op> FromIterator<O::Value> for LazySegtree<O> where O::Value: Clone, O::Operator: Clone, { fn from_iter<T: IntoIterator<Item = O::Value>>(iter: T) -> Self { Self::new(&iter.into_iter().collect::<Vec<_>>()) } } fn open<B: RangeBounds<usize>>(bounds: B, n: usize) -> (usize, usize) { use std::ops::Bound; let start = match bounds.start_bound() { Bound::Unbounded => 0, Bound::Included(&x) => x, Bound::Excluded(&x) => x + 1, }; let end = match bounds.end_bound() { Bound::Unbounded => n, Bound::Included(&x) => x + 1, Bound::Excluded(&x) => x, }; (start, end) } } // }}} // fp {{{ // https://ngtkana.github.io/ac-adapter-rs/fp/index.html #[allow(dead_code)] mod fp { mod ext_gcd { pub(crate) fn mod_inv<const P: u64>(x: u64) -> u64 { debug_assert!(P % 2 == 1); debug_assert!(P < 1 << 31); debug_assert!(x < P); mod_inv_signed(x as i64, P as i64) as u64 } fn mod_inv_signed(a: i64, m: i64) -> i64 { debug_assert!(a > 0); debug_assert!(m > 0); if a == 1 { return 1; } m + (1 - m * mod_inv_signed(m % a, a)) / a } } mod factorial { use super::Fp; use std::ops::Index; pub struct Factorial<const P: u64> { fact: Vec<Fp<P>>, inv_fact: Vec<Fp<P>>, } impl<const P: u64> Factorial<P> { pub fn new(length: usize) -> Self { let mut fact = vec![Fp::<P>::new(1); length + 1]; let mut inv_fact = vec![Fp::<P>::new(1); length + 1]; for i in 1..=length { fact[i] = fact[i - 1] * Fp::<P>::new(i as u64); } inv_fact[length] = fact[length].inv(); for i in (1..=length).rev() { inv_fact[i - 1] = inv_fact[i] * Fp::<P>::new(i as u64); } Self { fact, inv_fact } } pub fn fact(&self, n: usize) -> Fp<P> { self.fact[n] } pub fn inv_fact(&self, n: usize) -> Fp<P> { self.inv_fact[n] } pub fn perm(&self, n: usize, k: usize) -> Fp<P> { self.fact[n] * self.inv_fact[n - k] } pub fn comb(&self, n: usize, k: usize) -> Fp<P> { self.fact[n] * self.inv_fact[n - k] * self.inv_fact[k] } pub fn binom(&self, n: usize, k: usize) -> Fp<P> { self.comb(n, k) } pub fn comb_or_zero(&self, n: usize, k: isize) -> Fp<P> { if k < 0 || k as usize > n { Fp::<P>::new(0) } else { self.comb(n, k as usize) } } pub fn comb_with_reputation(&self, n: usize, k: usize) -> Fp<P> { assert!(n > 0 || k > 0); self.comb(n + k - 1, k) } } impl<const P: u64> Index<usize> for Factorial<P> { type Output = Fp<P>; fn index(&self, index: usize) -> &Self::Output { &self.fact[index] } } } mod fourier { use super::mod_inv; use super::Fp; use super::PrimitiveRoot; const P1: u64 = 924844033; const P2: u64 = 998244353; const P3: u64 = 1012924417; type F1 = Fp<P1>; type F2 = Fp<P2>; type F3 = Fp<P3>; pub fn fps_mul<const P: u64>(a: impl AsRef<[Fp<P>]>, b: impl AsRef<[Fp<P>]>) -> Vec<Fp<P>> where (): PrimitiveRoot<P>, { let a = a.as_ref(); let b = b.as_ref(); if a.is_empty() || b.is_empty() { return vec![]; } let mut a = a.to_vec(); let mut b = b.to_vec(); let n = a.len() + b.len() - 1; let len = n.next_power_of_two(); a.resize(len, Fp::new(0)); b.resize(len, Fp::new(0)); fft(&mut a); fft(&mut b); for (a, b) in a.iter_mut().zip(b.iter()) { *a *= *b; } ifft(&mut a); a.truncate(n); a } pub fn any_mod_fps_mul<const P: u64>(a: &[Fp<P>], b: &[Fp<P>]) -> Vec<Fp<P>> { let v1 = fps_mul( a.iter().map(|&x| F1::new(x.value())).collect::<Vec<_>>(), b.iter().map(|&x| F1::new(x.value())).collect::<Vec<_>>(), ); let v2 = fps_mul( a.iter().map(|&x| F2::new(x.value())).collect::<Vec<_>>(), b.iter().map(|&x| F2::new(x.value())).collect::<Vec<_>>(), ); let v3 = fps_mul( a.iter().map(|&x| F3::new(x.value())).collect::<Vec<_>>(), b.iter().map(|&x| F3::new(x.value())).collect::<Vec<_>>(), ); v1.into_iter() .zip(v2) .zip(v3) .map(|((e1, e2), e3)| garner(e1, e2, e3)) .collect::<Vec<_>>() } pub fn fft<const P: u64>(f: &mut [Fp<P>]) where (): PrimitiveRoot<P>, { let n = f.len(); assert!(n.is_power_of_two()); assert!((P - 1) % n as u64 == 0); let mut root = <() as PrimitiveRoot<P>>::VALUE.pow((P - 1) / f.len() as u64); let fourth = <() as PrimitiveRoot<P>>::VALUE.pow((P - 1) / 4); let mut fft_len = n; while 4 <= fft_len { let quarter = fft_len / 4; for f in f.chunks_mut(fft_len) { let mut c = Fp::new(1); for (((i, j), k), l) in (0..) .zip(quarter..) .zip(quarter * 2..) .zip(quarter * 3..) .take(quarter) { let c2 = c * c; let x = f[i] + f[k]; let y = f[j] + f[l]; let z = f[i] - f[k]; let w = fourth * (f[j] - f[l]); f[i] = x + y; f[j] = c2 * (x - y); f[k] = c * (z + w); f[l] = c2 * c * (z - w); c *= root; } } root *= root; root *= root; fft_len = quarter; } if fft_len == 2 { for f in f.chunks_mut(2) { let x = f[0]; let y = f[1]; f[0] = x + y; f[1] = x - y; } } } pub fn ifft<const P: u64>(f: &mut [Fp<P>]) where (): PrimitiveRoot<P>, { let n = f.len(); assert!(n.is_power_of_two()); let root = <() as PrimitiveRoot<P>>::VALUE.pow((P - 1) / f.len() as u64); let mut roots = std::iter::successors(Some(root.inv()), |x| Some(x * x)) .take(n.trailing_zeros() as usize + 1) .collect::<Vec<_>>(); roots.reverse(); let fourth = <() as PrimitiveRoot<P>>::VALUE.pow((P - 1) / 4).inv(); let mut quarter = 1_usize; if n.trailing_zeros() % 2 == 1 { for f in f.chunks_mut(2) { let x = f[0]; let y = f[1]; f[0] = x + y; f[1] = x - y; } quarter = 2; } while quarter != n { let fft_len = quarter * 4; let root = roots[fft_len.trailing_zeros() as usize]; for f in f.chunks_mut(fft_len) { let mut c = Fp::new(1); for (((i, j), k), l) in (0..) .zip(quarter..) .zip(quarter * 2..) .zip(quarter * 3..) .take(quarter) { let c2 = c * c; let x = f[i] + c2 * f[j]; let y = f[i] - c2 * f[j]; let z = c * (f[k] + c2 * f[l]); let w = fourth * c * (f[k] - c2 * f[l]); f[i] = x + z; f[j] = y + w; f[k] = x - z; f[l] = y - w; c *= root; } } quarter = fft_len; } let d = Fp::from(f.len()).inv(); f.iter_mut().for_each(|x| *x *= d); } fn garner<const P: u64>(x1: Fp<P1>, x2: Fp<P2>, x3: Fp<P3>) -> Fp<P> { let (x1, x2, x3) = (x1.value(), x2.value(), x3.value()); let x2 = ((x2 + (P2 - x1)) * mod_inv::<P2>(P1)) % P2; let x3 = (((x3 + (P3 - x1)) * mod_inv::<P3>(P1) % P3 + (P3 - x2)) * mod_inv::<P3>(P2)) % P3; Fp::new(x1 + P1 * (x2 + P2 * x3 % P)) } } use ext_gcd::mod_inv; pub use factorial::Factorial; pub use fourier::any_mod_fps_mul; pub use fourier::fft; pub use fourier::fps_mul; pub use fourier::ifft; use std::iter::Product; use std::iter::Sum; use std::mem::swap; use std::ops::Add; use std::ops::AddAssign; use std::ops::Div; use std::ops::DivAssign; use std::ops::Mul; use std::ops::MulAssign; use std::ops::Neg; use std::ops::Sub; use std::ops::SubAssign; #[macro_export] macro_rules! fp { ($value:expr) => { $crate::fp::Fp::from($value) }; ($value:expr; mod $p:expr) => { $crate::fp::Fp::<$p>::from($value) }; } pub trait PrimitiveRoot<const P: u64> { const VALUE: Fp<P>; } impl PrimitiveRoot<998244353> for () { const VALUE: Fp<998244353> = Fp::new(3); } impl PrimitiveRoot<1012924417> for () { const VALUE: Fp<1012924417> = Fp::new(5); } impl PrimitiveRoot<924844033> for () { const VALUE: Fp<924844033> = Fp::new(5); } #[derive(Clone, Copy, PartialEq, Eq, Hash)] pub struct Fp<const P: u64> { value: u64, } impl<const P: u64> Fp<P> { pub const fn new(value: u64) -> Self { Self { value: value % P } } pub const fn value(self) -> u64 { self.value } pub fn inv(self) -> Self { Self { value: mod_inv::<P>(self.value), } } pub fn pow(self, mut exp: u64) -> Self { let mut result = Self::new(1); let mut base = self; while exp > 0 { if exp & 1 == 1 { result *= base; } base *= base; exp >>= 1; } result } pub fn sign(pow: usize) -> Self { Self::new(if pow % 2 == 0 { 1 } else { P - 1 }) } } impl<const P: u64> std::fmt::Debug for Fp<P> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { pub fn berlekamp_massey_fp(a: i64, p: i64) -> [i64; 2] { let mut u0 = 0_i64; let mut v0 = 1_i64; let mut w0 = a * u0 + p * v0; let mut u1 = 1_i64; let mut v1 = 0_i64; let mut w1 = a * u1 + p * v1; while p <= w0 * w0 { let q = w0 / w1; u0 -= q * u1; v0 -= q * v1; w0 -= q * w1; swap(&mut u0, &mut u1); swap(&mut v0, &mut v1); swap(&mut w0, &mut w1); } [w0, u0] } if self.value == 0 { return write!(f, "0"); } let [mut num, mut den] = berlekamp_massey_fp(self.value as i64, P as i64); if den < 0 { num = -num; den = -den; } if den == 1 { write!(f, "{}", num) } else { write!(f, "{}/{}", num, den) } } } impl<const P: u64> std::fmt::Display for Fp<P> { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.value()) } } macro_rules! impl_from_signed { ($($t:ty),*) => { $( impl<const P: u64> From<$t> for Fp<P> { fn from(x: $t) -> Self { if x < 0 { -Self::new((P as i64 - x as i64) as u64) } else { Self::new(x as u64) } } } )* }; } impl_from_signed!(i8, i16, i32, i64, i128, isize); macro_rules! impl_from_unsigned { ($($t:ty),*) => { $( impl<const P: u64> From<$t> for Fp<P> { fn from(x: $t) -> Self { Self::new(x as u64) } } )* }; } impl_from_unsigned!(u8, u16, u32, u64, u128, usize); impl<const P: u64> AddAssign<Fp<P>> for Fp<P> { fn add_assign(&mut self, rhs: Fp<P>) { self.value += rhs.value; if self.value >= P { self.value -= P; } } } impl<const P: u64> SubAssign<Fp<P>> for Fp<P> { fn sub_assign(&mut self, rhs: Fp<P>) { if self.value < rhs.value { self.value += P; } self.value -= rhs.value; } } impl<const P: u64> MulAssign<Fp<P>> for Fp<P> { fn mul_assign(&mut self, rhs: Fp<P>) { self.value = self.value * rhs.value % P; } } #[allow(clippy::suspicious_op_assign_impl)] impl<const P: u64> DivAssign<Fp<P>> for Fp<P> { fn div_assign(&mut self, rhs: Fp<P>) { *self *= rhs.inv() } } macro_rules! fp_forward_ops { ($( $trait:ident, $trait_assign:ident, $fn:ident, $fn_assign:ident, )*) => {$( impl<const P: u64> $trait_assign<&Fp<P>> for Fp<P> { fn $fn_assign(&mut self, rhs: &Fp<P>) { self.$fn_assign(*rhs); } } impl<const P: u64, T: Into<Fp<P>>> $trait<T> for Fp<P> { type Output = Fp<P>; fn $fn(mut self, rhs: T) -> Self::Output { self.$fn_assign(rhs.into()); self } } impl<const P: u64> $trait<&Fp<P>> for Fp<P> { type Output = Fp<P>; fn $fn(self, rhs: &Fp<P>) -> Self::Output { self.$fn(*rhs) } } impl<const P: u64, T: Into<Fp<P>>> $trait<T> for &Fp<P> { type Output = Fp<P>; fn $fn(self, rhs: T) -> Self::Output { (*self).$fn(rhs.into()) } } impl<const P: u64> $trait<&Fp<P>> for &Fp<P> { type Output = Fp<P>; fn $fn(self, rhs: &Fp<P>) -> Self::Output { (*self).$fn(*rhs) } } )*}; } fp_forward_ops! { Add, AddAssign, add, add_assign, Sub, SubAssign, sub, sub_assign, Mul, MulAssign, mul, mul_assign, Div, DivAssign, div, div_assign, } impl<const P: u64> Neg for Fp<P> { type Output = Fp<P>; fn neg(mut self) -> Self::Output { if self.value > 0 { self.value = P - self.value; } self } } impl<const P: u64> Sum for Fp<P> { fn sum<I: Iterator<Item = Self>>(iter: I) -> Self { iter.fold(Self::new(0), |acc, x| acc + x) } } impl<'a, const P: u64> Sum<&'a Self> for Fp<P> { fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self { iter.copied().sum() } } impl<const P: u64> Product for Fp<P> { fn product<I: Iterator<Item = Self>>(iter: I) -> Self { iter.fold(Self::new(1), |acc, x| acc * x) } } impl<'a, const P: u64> Product<&'a Self> for Fp<P> { fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self { iter.copied().product() } } } // }}} // lg {{{ // https://ngtkana.github.io/ac-adapter-rs/lg/index.html #[allow(dead_code)] mod lg { use std::borrow::Borrow; use std::fmt; use std::iter::once; #[macro_export] macro_rules! lg { (@contents $head:expr $(, $tail:expr)*) => {{ $crate::__lg_internal!($head); $( eprint!(","); $crate::__lg_internal!($tail); )* eprintln!(); }}; ($($expr:expr),* $(,)?) => {{ eprint!("{}\u{276f}", line!()); $crate::lg!(@contents $($expr),*) }}; } #[doc(hidden)] #[macro_export] macro_rules! __lg_internal { ($value:expr) => {{ match $value { head => { eprint!( " {} = {}", stringify!($value), $crate::lg::__quiet(format!("{:?}", &head)) ); } } }}; } #[macro_export] macro_rules! dict { ($($(@field $field:ident)? $values:expr ),* $(,)?) => {{ #![allow(unused_assignments)] let mut dict = $crate::lg::Dict::default(); $( { let mut name = stringify!($values).to_string(); if name.starts_with("&") { name = name[1..].to_string(); } $( let name = format!("{name}.{field}", field = stringify!($field)); )? let values = $values; $( let values = values.iter().map(|value| &value.$field).collect::<Vec<_>>(); )? dict.add_row(name, values); } )* eprintln!("{}", dict); }}; } #[macro_export] macro_rules! rows { { $index_label:literal, $(@offset $offset:expr,)? $(@verticalbar $verticalbar:expr,)* $($(@$label:literal =>)? $values:expr),* $(,)? } => {{ #![allow(unused_assignments)] let mut rows = $crate::lg::Rows::default(); rows.line_number(line!()); $(rows.offset($offset);)? $(rows.verticalbar($verticalbar);)* rows.index_label($index_label); $({ let mut label = stringify!($values).to_string(); if label.starts_with("&") { label = label[1..].to_string(); } $({ let label_: &'static str = $label; label = label_.to_string(); })? rows.row(label, $values); })* eprintln!("{}", rows.to_string_table()); }}; } #[macro_export] macro_rules! table { { $(@$name:literal => )? $values:expr $(,)? } => {{ #![allow(unused_assignments)] let mut name = stringify!($values).to_string(); if name.starts_with("&") { name = name[1..].to_string(); } $({ let name_: &'static str = $name; name = name_.to_string(); })? let mut rows = $crate::lg::Rows::default(); rows.line_number(line!()); rows.table_name(name); #[allow(array_into_iter)] for (i, row) in $values.into_iter().enumerate() { rows.row(i.to_string(), row); } eprintln!("{}", rows.to_string_table()); }}; } #[doc(hidden)] pub fn __quiet(s: impl AsRef<str>) -> String { s.as_ref() .replace("340282366920938463463374607431768211455", "*") // u128 .replace("170141183460469231731687303715884105727", "*") // i128 .replace("18446744073709551615", "*") // u64 .replace("9223372036854775807", "*") // i64 .replace("-9223372036854775808", "*") // i64 .replace("4294967295", "*") // u32 .replace("2147483647", "*") // i32 .replace("-2147483648", "*") // i32 .replace("None", "*") .replace("Some", "") .replace("true", "#") .replace("false", ".") .replace(['"', '\''], "") } #[doc(hidden)] #[derive(Default)] pub struct Rows { line_number: String, index_label: String, offset: usize, verticalbars: Vec<usize>, table_name: String, rows: Vec<Row>, } impl Rows { pub fn line_number(&mut self, line_number: u32) -> &mut Self { self.line_number = format!("{}", line_number); self } pub fn index_label(&mut self, index_label: impl Into<String>) -> &mut Self { self.index_label = index_label.into(); self } pub fn offset(&mut self, offset: usize) -> &mut Self { self.offset = offset; self } pub fn verticalbar(&mut self, verticalbar: impl IntoIterator<Item = usize>) -> &mut Self { self.verticalbars.extend(verticalbar); self } pub fn table_name(&mut self, table_name: impl Into<String>) -> &mut Self { self.table_name = table_name.into(); self } pub fn row( &mut self, label: impl Into<String>, values: impl IntoIterator<Item = impl fmt::Debug>, ) -> &mut Self { self.rows.push(Row { label: label.into(), values: values .into_iter() .map(|value| __quiet(format!("{:?}", value))) .collect(), }); self } pub fn to_string_table(self) -> StringTable { let Self { line_number, index_label, offset, verticalbars, table_name, rows, } = self; let w = rows .iter() .map(|row| row.values.len()) .max() .unwrap_or_default(); let mut verticalbar_count = vec![0; w + 1]; for &v in &verticalbars { if (offset..=offset + w).contains(&v) { verticalbar_count[v - offset] += 1; } } StringTable { head: StringRow { label: format!( "{line_number}❯ {table_name}{index_label}", index_label = if index_label.is_empty() { String::new() } else { format!("[{}]", index_label) } ), values: (offset..offset + w) .map(|index| index.to_string()) .collect(), }, body: rows .iter() .map(|row| StringRow { label: row.label.clone(), values: row.values.clone(), }) .collect(), verticalbar_count, } } } struct Row { label: String, values: Vec<String>, } #[doc(hidden)] pub struct StringTable { head: StringRow, body: Vec<StringRow>, verticalbar_count: Vec<usize>, } impl fmt::Display for StringTable { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let Self { head, body, verticalbar_count, } = self; let w = body .iter() .map(|row| row.values.len()) .max() .unwrap_or_default(); let label_width = once(head.label.chars().count()) .chain(body.iter().map(|row| row.label.chars().count())) .max() .unwrap(); let value_width = (0..w) .map(|j| { once(j.to_string().len()) .chain( body.iter() .map(|row| row.values.get(j).map_or(0, |s| s.chars().count())), ) .max() .unwrap() }) .collect::<Vec<_>>(); // Heading gray(f)?; write!( f, "{}", head.to_string(label_width, &value_width, verticalbar_count, true) )?; resetln(f)?; // Body for row in body { write!( f, "{}", row.to_string(label_width, &value_width, verticalbar_count, false) )?; writeln!(f)?; } Ok(()) } } struct StringRow { label: String, values: Vec<String>, } impl StringRow { fn to_string( &self, label_width: usize, value_width: &[usize], varticalbars_count: &[usize], label_align_left: bool, ) -> String { let Self { label, values } = self; let w = value_width.len(); let mut s = String::new(); s.push_str(&if label_align_left { format!("{label:<label_width$} |") } else { format!("{label:^label_width$} |") }); for j in 0..w { let value_width = value_width[j]; s.push_str("|".repeat(varticalbars_count[j]).as_str()); if varticalbars_count[j] == 0 && j != 0 && value_width <= 1 { s.push(' '); } match values.get(j) { Some(value) => { s.push_str(&format!(" {value:>value_width$}",)); } None => { s.push_str(" ".repeat(value_width + 1).as_str()); } } } s } } const GRAY: &str = "\x1b[48;2;127;127;127;37m"; const RESET: &str = "\x1b[0m"; fn gray(f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "{GRAY}") } fn resetln(f: &mut fmt::Formatter<'_>) -> fmt::Result { writeln!(f, "{RESET}") } pub fn bools<B, I>(iter: I) -> String where B: Borrow<bool>, I: IntoIterator<Item = B>, { format!( "[{}]", iter.into_iter() .map(|b| ['.', '#'][usize::from(*(b.borrow()))]) .collect::<String>(), ) } #[derive(Default)] pub struct Dict { dict: Vec<(String, Vec<String>)>, } impl Dict { pub fn add_row<T: std::fmt::Debug>( &mut self, key: impl AsRef<str>, values: impl IntoIterator<Item = T>, ) { self.dict.push(( key.as_ref().to_string(), values .into_iter() .map(|value| format!("{:?}", value)) .collect(), )); } } impl std::fmt::Display for Dict { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let max_nvalues = self .dict .iter() .map(|(_, values)| values.len()) .max() .unwrap_or(0); let key_width = self .dict .iter() .map(|(key, _)| key.len()) .max() .unwrap_or(0); let mut value_widths = vec![0; max_nvalues]; for (_, values) in &self.dict { for (i, value) in values.iter().enumerate() { value_widths[i] = value_widths[i].max(value.len()); } } write!(f, "{GRAY} {key:key_width$} |", key = "")?; for (i, &value_width) in value_widths.iter().enumerate() { write!(f, " {i:value_width$} ")?; } writeln!(f, "{RESET}")?; for (key, values) in &self.dict { write!(f, " {key:key_width$} |")?; for (value, &value_width) in values.iter().zip(&value_widths) { write!(f, " {value:value_width$} ",)?; } writeln!(f)?; } Ok(()) } } } // }}}