結果
| 問題 |
No.650 行列木クエリ
|
| コンテスト | |
| ユーザー |
 ngtkana
|
| 提出日時 | 2024-06-24 14:47:03 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
AC
|
| 実行時間 | 72 ms / 2,000 ms |
| コード長 | 40,734 bytes |
| コンパイル時間 | 14,127 ms |
| コンパイル使用メモリ | 387,260 KB |
| 実行使用メモリ | 33,172 KB |
| 最終ジャッジ日時 | 2024-06-24 14:47:20 |
| 合計ジャッジ時間 | 15,429 ms |
|
ジャッジサーバーID （参考情報） |
judge5 / judge1 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 1 |
| other | AC * 10 |
コンパイルメッセージ
warning: unused import: `factorial::Factorial`
--> src/main.rs:1046:13
|
1046 | pub use factorial::Factorial;
| ^^^^^^^^^^^^^^^^^^^^
|
= note: `#[warn(unused_imports)]` on by default
warning: unused import: `fourier::any_mod_fps_mul`
--> src/main.rs:1047:13
|
1047 | pub use fourier::any_mod_fps_mul;
| ^^^^^^^^^^^^^^^^^^^^^^^^
warning: unused import: `fourier::fft`
--> src/main.rs:1048:13
|
1048 | pub use fourier::fft;
| ^^^^^^^^^^^^
warning: unused import: `fourier::fps_mul`
--> src/main.rs:1049:13
|
1049 | pub use fourier::fps_mul;
| ^^^^^^^^^^^^^^^^
warning: unused import: `fourier::ifft`
--> src/main.rs:1050:13
|
1050 | pub use fourier::ifft;
| ^^^^^^^^^^^^^
ソースコード
use proconio::input;
use segtree::Segtree;
use std::iter::FusedIterator;
type Fp = fp::Fp<1_000_000_007>;
fn main() {
input! {
n: usize,
edges: [(usize, usize); n - 1],
q: usize,
}
let (hld, _g) = Hld::from_edges(0, &edges);
let Hld { index, .. } = &hld;
let mut segtree = Segtree::<O>::from_iter(vec![<O as segtree::Op>::identity(); n]);
for _ in 0..q {
input! {
com: String,
}
match com.as_str() {
"x" => {
input! {
e: usize,
x00: u64,
x01: u64,
x10: u64,
x11: u64,
}
let (i, j) = edges[e];
*segtree.entry(index[i].max(index[j])) =
[[fp!(x00), fp!(x01)], [fp!(x10), fp!(x11)]];
}
"g" => {
input! {
i: usize,
j: usize,
}
let mut ans = MATRIX_IDENTITY;
for (l, r, last, reversed) in hld.path_segments_by_index(i, j) {
assert!(!reversed);
ans = matrix_mul(segtree.fold(l + usize::from(last)..=r), ans);
}
let [[x00, x01], [x10, x11]] = ans;
println!("{x00} {x01} {x10} {x11}");
}
_ => unreachable!(),
}
}
}
type Matrix = [[Fp; 2]; 2];
const MATRIX_IDENTITY: Matrix = [[Fp::new(1), Fp::new(0)], [Fp::new(0), Fp::new(1)]];
fn matrix_mul(lhs: Matrix, rhs: Matrix) -> Matrix {
let mut result = [[fp!(0); 2]; 2];
for i in 0..2 {
for j in 0..2 {
for k in 0..2 {
result[i][k] += lhs[i][j] * rhs[j][k];
}
}
}
result
}
enum O {}
impl segtree::Op for O {
type Value = Matrix;
fn identity() -> Self::Value {
MATRIX_IDENTITY
}
fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value {
matrix_mul(*lhs, *rhs)
}
}
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>>) {
Self::from_edge_iterator(root, edges.iter().copied())
}
pub fn from_edge_iterator(
root: usize,
edges: impl ExactSizeIterator<Item = (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 path_segments(&self, from: usize, to: usize) -> PathSegments<'_> {
PathSegments {
hld: self,
from,
to,
exhausted: false,
}
}
pub fn path_segments_by_index(
&self,
i: usize,
j: usize,
) -> impl Iterator<Item = (usize, usize, bool, bool)> + '_ {
self.path_segments(i, j)
.map(|(i, j, last, reversed)| (self.index[i], self.index[j], last, reversed))
}
pub fn ledacy_iter_v(&self, i: usize, j: usize) -> impl Iterator<Item = (usize, usize)> + '_ {
self.path_segments_by_index(i, j)
.map(|(i, j, _last, _revresed)| (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 PathSegments<'a> {
hld: &'a Hld,
from: usize,
to: usize,
exhausted: bool,
}
impl Iterator for PathSegments<'_> {
// i: usize root-side,
// j: usize leaf-side,
// last: bool contains-lca
// reversed: bool false if from --> to and i --> j coincide
type Item = (usize, usize, bool, bool);
fn next(&mut self) -> Option<Self::Item> {
(!self.exhausted).then_some(())?;
let Self {
hld,
from: i,
to: j,
..
} = *self;
let Hld {
index,
head,
parent,
} = hld;
Some(if head[i] == head[j] {
self.exhausted = true;
if index[i] <= index[j] {
(i, j, true, false)
} else {
(j, i, true, true)
}
} else {
if index[i] < index[j] {
self.to = parent[head[j]];
(head[j], j, false, false)
} else {
self.from = parent[head[i]];
(head[i], i, false, true)
}
})
}
}
impl FusedIterator for PathSegments<'_> {}
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
}
// segtree {{{
// https://ngtkana.github.io/ac-adapter-rs/segtree/index.html
#[allow(dead_code)]
mod segtree {
use core::fmt;
use std::collections::BTreeMap;
use std::iter::FromIterator;
use std::ops::Deref;
use std::ops::DerefMut;
use std::ops::Index;
use std::ops::RangeBounds;
pub trait Op {
type Value;
fn identity() -> Self::Value;
fn op(lhs: &Self::Value, rhs: &Self::Value) -> Self::Value;
}
pub struct Segtree<O: Op> {
values: Vec<O::Value>,
}
impl<O: Op> Segtree<O> {
pub fn from_len(n: usize) -> Self
where
O::Value: Clone,
{
Self::new(&vec![O::identity(); n])
}
pub fn new(values: &[O::Value]) -> Self
where
O::Value: 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 }
}
pub fn fold<R: RangeBounds<usize>>(&self, range: R) -> O::Value {
let n = self.values.len() / 2;
let (mut start, mut end) = open(range, n);
assert!(start <= end && end <= n);
start += n;
end += n;
let mut left = O::identity();
let mut right = O::identity();
while start < end {
if start % 2 == 1 {
left = O::op(&left, &self.values[start]);
start += 1;
}
if end % 2 == 1 {
end -= 1;
right = O::op(&self.values[end], &right);
}
start /= 2;
end /= 2;
}
O::op(&left, &right)
}
pub fn entry(&mut self, index: usize) -> Entry<O> {
let n = self.values.len() / 2;
Entry {
segtree: self,
index: n + index,
}
}
pub fn iter(&self) -> impl Iterator<Item = &O::Value> {
self.values[self.values.len() / 2..].iter()
}
pub fn as_slice(&self) -> &[O::Value] {
&self.values[self.values.len() / 2..]
}
}
impl<O: Op> fmt::Debug for Segtree<O>
where
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Segtree")
.field("values", &self.values)
.finish()
}
}
impl<O: Op> FromIterator<O::Value> for Segtree<O>
where
O::Value: Clone,
{
fn from_iter<I: IntoIterator<Item = O::Value>>(iter: I) -> Self {
Self::new(&iter.into_iter().collect::<Vec<_>>())
}
}
impl<O: Op> Index<usize> for Segtree<O> {
type Output = O::Value;
fn index(&self, index: usize) -> &Self::Output {
&self.values[self.values.len() / 2 + index]
}
}
pub struct Entry<'a, O: Op> {
segtree: &'a mut Segtree<O>,
index: usize,
}
impl<'a, O: Op> Drop for Entry<'a, O> {
fn drop(&mut self) {
let mut index = self.index;
while index != 0 {
index /= 2;
self.segtree.values[index] = O::op(
&self.segtree.values[index * 2],
&self.segtree.values[index * 2 + 1],
);
}
}
}
impl<'a, O: Op> Deref for Entry<'a, O> {
type Target = O::Value;
fn deref(&self) -> &Self::Target {
&self.segtree.values[self.index]
}
}
impl<'a, O: Op> DerefMut for Entry<'a, O> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.segtree.values[self.index]
}
}
impl<'a, O: Op> fmt::Debug for Entry<'a, O>
where
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Entry").field("index", &self.index).finish()
}
}
pub struct SparseSegtree<K, O: Op> {
inner: Segtree<O>,
keys: Vec<K>,
}
impl<K: Ord, O: Op> SparseSegtree<K, O> {
pub fn new(kv: &[(K, O::Value)]) -> Self
where
K: Clone,
O::Value: Clone,
{
let mut keys = kv.iter().map(|(k, _)| k.clone()).collect::<Vec<_>>();
keys.sort();
let values = kv.iter().map(|(_, v)| v.clone()).collect::<Vec<_>>();
Self {
inner: Segtree::new(&values),
keys,
}
}
pub fn fold<R: RangeBounds<K>>(&self, range: R) -> O::Value {
let (start, end) = open_key(range, &self.keys);
self.inner.fold(start..end)
}
pub fn entry(&mut self, key: &K) -> Entry<'_, O> {
let index = self.keys.binary_search(key).unwrap() + self.keys.len();
Entry {
segtree: &mut self.inner,
index,
}
}
pub fn keys(&self) -> &[K] {
&self.keys
}
pub fn iter(&self) -> impl Iterator<Item = (&K, &O::Value)> {
self.keys.iter().zip(self.inner.as_slice())
}
pub fn collect_map(&self) -> BTreeMap<K, O::Value>
where
K: Clone,
O::Value: Clone,
{
self.keys
.iter()
.cloned()
.zip(self.inner.iter().cloned())
.collect()
}
}
impl<K, O: Op> fmt::Debug for SparseSegtree<K, O>
where
K: fmt::Debug,
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SparseSegtree")
.field("inner", &self.inner)
.field("keys", &self.keys)
.finish()
}
}
impl<K: Ord, O: Op> FromIterator<(K, O::Value)> for SparseSegtree<K, O>
where
K: Clone,
O::Value: Clone,
{
fn from_iter<I: IntoIterator<Item = (K, O::Value)>>(iter: I) -> Self {
Self::new(&iter.into_iter().collect::<Vec<_>>())
}
}
impl<K: Ord, O: Op> Index<K> for SparseSegtree<K, O> {
type Output = O::Value;
fn index(&self, key: K) -> &Self::Output {
&self.inner[self.keys.binary_search(&key).unwrap()]
}
}
pub struct Sparse2dSegtree<K, L, O: Op> {
segtrees: Vec<SparseSegtree<L, O>>,
keys: Vec<K>,
}
impl<K, L, O: Op> Sparse2dSegtree<K, L, O>
where
K: Ord + Clone,
L: Ord + Clone,
O::Value: Clone,
{
pub fn new(points: &[(K, L, O::Value)]) -> Self {
let mut keys = points.iter().map(|(k, _, _)| k.clone()).collect::<Vec<_>>();
keys.sort();
keys.dedup();
let mut lvs = vec![vec![]; keys.len() * 2];
for (k, l, v) in points {
let mut i = keys.binary_search(k).unwrap();
i += keys.len();
while i != 0 {
lvs[i].push((l.clone(), v.clone()));
i /= 2;
}
}
let segtrees = lvs
.into_iter()
.map(|lvs_| {
let mut ls = lvs_.iter().map(|(l, _)| l).collect::<Vec<_>>();
ls.sort();
ls.dedup();
let mut lvs = ls
.iter()
.map(|&l| (l.clone(), O::identity()))
.collect::<Vec<_>>();
for (l, v) in &lvs_ {
let i = ls.binary_search(&l).unwrap();
lvs[i].1 = O::op(&lvs[i].1, v);
}
SparseSegtree::new(&lvs)
})
.collect::<Vec<_>>();
Self { segtrees, keys }
}
pub fn fold(&self, i: impl RangeBounds<K>, j: impl RangeBounds<L> + Clone) -> O::Value {
let (mut i0, mut i1) = open_key(i, &self.keys);
i0 += self.keys.len();
i1 += self.keys.len();
let mut left = O::identity();
let mut right = O::identity();
while i0 < i1 {
if i0 % 2 == 1 {
left = O::op(&left, &self.segtrees[i0].fold(j.clone()));
i0 += 1;
}
if i1 % 2 == 1 {
i1 -= 1;
right = O::op(&self.segtrees[i1].fold(j.clone()), &right);
}
i0 /= 2;
i1 /= 2;
}
O::op(&left, &right)
}
pub fn apply(&mut self, k: &K, l: &L, mut f: impl FnMut(&mut O::Value)) {
let mut i = self.keys.binary_search(k).unwrap();
i += self.keys.len();
while i != 0 {
f(&mut self.segtrees[i].entry(l));
i /= 2;
}
}
pub fn iter(&self) -> impl Iterator<Item = (&K, &L, &O::Value)> {
self.keys
.iter()
.zip(self.segtrees[self.keys.len()..].iter())
.flat_map(|(k, segtree)| segtree.iter().map(move |(l, v)| (k, l, v)))
}
pub fn collect_map(&self) -> BTreeMap<(K, L), O::Value>
where
K: Clone,
L: Clone,
O::Value: Clone,
{
self.keys
.iter()
.flat_map(|k| {
self.segtrees[self.keys.len() + self.keys.binary_search(k).unwrap()]
.iter()
.map(move |(l, v)| ((k.clone(), l.clone()), v.clone()))
})
.collect()
}
}
impl<K, L, O: Op> fmt::Debug for Sparse2dSegtree<K, L, O>
where
K: fmt::Debug,
L: fmt::Debug,
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Sparse2dSegtree")
.field("segtrees", &self.segtrees)
.field("keys", &self.keys)
.finish()
}
}
impl<K, L, O: Op> FromIterator<(K, L, O::Value)> for Sparse2dSegtree<K, L, O>
where
K: Ord + Clone,
L: Ord + Clone,
O::Value: Clone,
{
fn from_iter<I: IntoIterator<Item = (K, L, O::Value)>>(iter: I) -> Self {
Self::new(&iter.into_iter().collect::<Vec<_>>())
}
}
impl<K: Ord, L: Ord, O: Op> Index<K> for Sparse2dSegtree<K, L, O> {
type Output = SparseSegtree<L, O>;
fn index(&self, i: K) -> &Self::Output {
&self.segtrees[self.keys.binary_search(&i).unwrap() + self.keys.len()]
}
}
impl<K: Ord, L: Ord, O: Op> Index<(K, L)> for Sparse2dSegtree<K, L, O> {
type Output = O::Value;
fn index(&self, (i, j): (K, L)) -> &Self::Output {
&self.segtrees[self.keys.binary_search(&i).unwrap() + self.keys.len()][j]
}
}
pub struct Dense2dSegtree<O: Op> {
values: Vec<Vec<O::Value>>,
}
impl<O: Op> Dense2dSegtree<O> {
pub fn new(values: &[Vec<O::Value>]) -> Self
where
O::Value: Clone,
{
let values_ = values;
let h = values.len();
let w = values.get(0).map_or(0, Vec::len);
let mut values = vec![vec![O::identity(); 2 * w]; 2 * h];
for (values, values_) in values[h..].iter_mut().zip(values_) {
values[w..].clone_from_slice(values_);
for j in (1..w).rev() {
values[j] = O::op(&values[j * 2], &values[j * 2 + 1]);
}
}
for i in (1..h).rev() {
for j in 0..2 * w {
values[i][j] = O::op(&values[i * 2][j], &values[i * 2 + 1][j]);
}
}
Self { values }
}
pub fn fold(&self, i: impl RangeBounds<usize>, j: impl RangeBounds<usize>) -> O::Value {
let h = self.values.len() / 2;
let w = self.values.get(0).map_or(0, |v| v.len() / 2);
let (mut i0, mut i1) = open(i, h);
assert!(i0 <= i1 && i1 <= h);
let (mut j0, mut j1) = open(j, w);
assert!(j0 <= j1 && j1 <= w);
i0 += h;
i1 += h;
j0 += w;
j1 += w;
let mut left = O::identity();
let mut right = O::identity();
while i0 < i1 {
if i0 % 2 == 1 {
let mut j0 = j0;
let mut j1 = j1;
while j0 < j1 {
if j0 % 2 == 1 {
left = O::op(&left, &self.values[i0][j0]);
j0 += 1;
}
if j1 % 2 == 1 {
j1 -= 1;
right = O::op(&self.values[i0][j1], &right);
}
j0 /= 2;
j1 /= 2;
}
i0 += 1;
}
if i1 % 2 == 1 {
i1 -= 1;
let mut j0 = j0;
let mut j1 = j1;
while j0 < j1 {
if j0 % 2 == 1 {
left = O::op(&left, &self.values[i1][j0]);
j0 += 1;
}
if j1 % 2 == 1 {
j1 -= 1;
right = O::op(&self.values[i1][j1], &right);
}
j0 /= 2;
j1 /= 2;
}
}
i0 /= 2;
i1 /= 2;
}
O::op(&left, &right)
}
pub fn entry(&mut self, i: usize, j: usize) -> Dense2dEntry<O> {
let h = self.values.len() / 2;
let w = self.values.get(0).map_or(0, |v| v.len() / 2);
Dense2dEntry {
segtree: self,
i: h + i,
j: w + j,
}
}
pub fn iter(&self) -> impl Iterator<Item = &[O::Value]> {
self.values[self.values.len() / 2..]
.iter()
.map(|v| &v[v.len() / 2..])
}
pub fn collect_vec(&self) -> Vec<Vec<O::Value>>
where
O::Value: Clone,
{
self.iter().map(<[_]>::to_vec).collect()
}
}
impl<O: Op> fmt::Debug for Dense2dSegtree<O>
where
O::Value: fmt::Debug,
{
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Dense2dSegtree")
.field("values", &self.values)
.finish()
}
}
impl<O: Op> Index<usize> for Dense2dSegtree<O> {
type Output = [O::Value];
fn index(&self, index: usize) -> &Self::Output {
&self.values[self.values.len() / 2 + index]
}
}
pub struct Dense2dEntry<'a, O: Op> {
segtree: &'a mut Dense2dSegtree<O>,
i: usize,
j: usize,
}
impl<'a, O: Op> Drop for Dense2dEntry<'a, O> {
fn drop(&mut self) {
let mut i = self.i;
let mut j = self.j / 2;
while j != 0 {
self.segtree.values[i][j] = O::op(
&self.segtree.values[i][2 * j],
&self.segtree.values[i][2 * j + 1],
);
j /= 2;
}
i /= 2;
while i != 0 {
let mut j = self.j;
while j != 0 {
self.segtree.values[i][j] = O::op(
&self.segtree.values[i * 2][j],
&self.segtree.values[i * 2 + 1][j],
);
j /= 2;
}
i /= 2;
}
}
}
impl<'a, O: Op> Deref for Dense2dEntry<'a, O> {
type Target = O::Value;
fn deref(&self) -> &Self::Target {
&self.segtree.values[self.i][self.j]
}
}
impl<'a, O: Op> DerefMut for Dense2dEntry<'a, O> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.segtree.values[self.i][self.j]
}
}
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)
}
fn open_key<K: Ord, B: RangeBounds<K>>(bounds: B, keys: &[K]) -> (usize, usize) {
use std::ops::Bound;
let start = match bounds.start_bound() {
Bound::Unbounded => 0,
Bound::Included(x) => match keys.binary_search(x) {
Ok(i) | Err(i) => i,
},
Bound::Excluded(x) => match keys.binary_search(x) {
Ok(i) => i + 1,
Err(i) => i,
},
};
let end = match bounds.end_bound() {
Bound::Unbounded => keys.len(),
Bound::Included(x) => match keys.binary_search(x) {
Ok(i) => i + 1,
Err(i) => i,
},
Bound::Excluded(x) => match keys.binary_search(x) {
Ok(i) | Err(i) => i,
},
};
(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()
}
}
}
// }}}