結果
| 問題 |
No.2801 Unique Maximum
|
| コンテスト | |
| ユーザー |
 akakimidori
|
| 提出日時 | 2025-11-10 00:49:03 |
| 言語 | Rust (1.83.0 + proconio) |
| 結果 |
AC
|
| 実行時間 | 1,148 ms / 4,000 ms |
| コード長 | 43,003 bytes |
| コンパイル時間 | 15,157 ms |
| コンパイル使用メモリ | 400,092 KB |
| 実行使用メモリ | 17,824 KB |
| 最終ジャッジ日時 | 2025-11-10 00:51:36 |
| 合計ジャッジ時間 | 27,551 ms |
|
ジャッジサーバーID （参考情報） |
judge3 / judge4 |
（要ログイン）
| ファイルパターン | 結果 |
|---|---|
| sample | AC * 3 |
| other | AC * 21 |
ソースコード
// F(x+x^2) = F + F^2
// F[..3] = [0, 1, M]
// [x^(n+1)] F が答え
// [x^i]F (i>=3) を計算したい
// i次は同じ
// i+1次
// i f vs 2f
// 解ける
fn main() {
input! {
n: usize,
m: usize,
}
let mut conv = OnlineConvolution::new();
let mut comp = OnlineSmallComposition::new(vec![M::zero(), M::one(), M::one()]);
let mut ans = vec![M::zero(); n + 2];
ans[1] = M::one();
ans[2] = M::from(m);
for i in 0..=(n + 1) {
if i > 2 {
let l = comp.find_assume(i + 1);
let r = conv.find_assume(i + 1);
ans[i] = (r - l) * M::from(i - 2).inv();
}
conv.next(ans[i], ans[i]);
comp.next(ans[i]);
}
println!("{}", ans[n + 1]);
}
type M = ModInt<998244353>;
// 合成する多項式は小さいことを想定している
// もうちょいマシな実装にしたい
#[derive(Debug)]
pub struct OnlineSmallComposition<T> {
f: Vec<T>,
h: Vec<T>,
stack: Vec<Vec<T>>,
pow: Vec<Vec<T>>,
pos: usize,
}
impl<T> OnlineSmallComposition<T>
where
T: Copy + Field + From<usize> + std::fmt::Debug,
[T]: ArrayConvolution<Item = T>,
{
pub fn new(mut g: Vec<T>) -> Self {
assert!(g.len() > 1 && g[0].is_zero() && !g[1].is_zero());
g.remove(0);
Self {
f: vec![],
h: vec![],
stack: vec![],
pow: vec![g],
pos: 0,
}
}
pub fn next(&mut self, f: T) -> T {
self.f.push(f);
if self.pos == 0 {
self.stack.push(vec![f]);
self.h.push(f);
self.pos += 1;
return f;
}
let x = self.pos;
let k = x.trailing_zeros();
let l = self
.stack
.last()
.unwrap()
.iter()
.take(2 << k)
.cloned()
.collect::<Vec<_>>();
let p = small_pow(self.pow[0].clone(), x - (1 << k), 2 << k);
let mut m = l.convolution(&p);
m.truncate(2 << k);
if 2 << k > self.h.len() {
self.h.resize(2 << k, T::zero());
}
self.h[self.pos..].add_assign(&m[1 << k..]);
self.h[self.pos] = self.h[self.pos] + f * pow(self.pow[0][0], self.pos);
self.stack.push(vec![f]);
let o = x.trailing_ones() as usize;
while self.pow.len() + 1 <= o {
let p = self.pow.last().unwrap();
let q = p.convolution(p);
self.pow.push(q);
}
for i in 0..o {
let a = self.stack.pop().unwrap().convolution(&self.pow[i]);
let mut b = self.stack.pop().unwrap();
b.resize((1 << i) + a.len(), T::zero());
b[(1 << i)..].add_assign(&a);
self.stack.push(b);
}
self.pos += 1;
self.h[self.pos - 1]
}
pub fn find_assume(&mut self, x: usize) -> T {
if x < self.pos {
return self.h[x];
}
if self.pos == 0 {
return T::zero();
}
let mut res = self.h.get(x).cloned().unwrap_or(T::zero());
let mut poly = vec![];
let mut pos = self.pos;
let mut top = self.stack.len();
for i in 0.. {
if pos >> i & 1 == 1 {
while self.pow.len() <= i {
let p = self.pow.last().unwrap();
let q = p.convolution(p);
self.pow.push(q);
}
poly = poly.convolution(&self.pow[i]);
poly.splice(0..0, (0..(1 << i)).map(|_| T::zero()));
poly.add_assign(&self.stack[top - 1]);
top -= 1;
pos -= 1 << i;
}
if poly.len() > 0 && pos + (2 << i) > x {
let pow = small_pow(self.pow[0].clone(), pos, x - pos + 1);
for i in 0..pow.len() {
if let Some(&v) = poly.get(x - pos - i) {
res = res + v * pow[i];
}
}
break;
}
}
res
}
}
// f^m の[0..n)を求める
pub fn small_pow<T>(mut f: Vec<T>, m: usize, mut n: usize) -> Vec<T>
where
T: Field + From<usize> + Copy,
{
let s = f.iter().position(|f| !f.is_zero());
if s.map_or(true, |s| s * m >= n) {
return vec![T::zero(); n];
}
let s = s.unwrap();
if s > 0 {
n -= m * s;
f.drain(..s);
}
let f0 = f[0];
let inv = T::one() / f0;
for f in f.iter_mut() {
*f = *f * inv;
}
let mut dp = vec![T::zero(); n + s * m];
dp[0] = pow(f0, m);
let pc = Precalc::new(n);
for i in 1..n {
let mut s = T::zero();
for (j, (f, dp)) in f[1..].iter().zip(dp[..i].iter().rev()).enumerate() {
s = s + *f * T::from(j + 1) * *dp;
}
s = s * T::from(m);
for (j, f) in f[1..].iter().enumerate().take(i) {
s = s - *f * T::from(i - 1 - j) * dp[i - 1 - j];
}
dp[i] = s * pc.inv(i);
}
dp.rotate_right(m * s);
dp
}
pub struct OnlineConvolution<T> {
f: Vec<T>,
g: Vec<T>,
h: Vec<T>,
pos: usize,
}
impl<T> OnlineConvolution<T>
where
T: Copy + Field,
[T]: ArrayConvolution<Item = T>,
{
pub fn new() -> Self {
Self {
f: vec![],
g: vec![],
h: vec![],
pos: 0,
}
}
pub fn next(&mut self, f: T, g: T) -> T {
self.f.push(f);
self.g.push(g);
let a = self.pos + 2;
let len = 1 << a.trailing_zeros();
if a == len {
let c = self.f.convolution(&self.g);
self.h.extend(c[self.pos..].iter().copied());
} else {
let r = self.f[self.pos + 1 - len..]
.iter()
.cloned()
.rev()
.collect::<Vec<_>>();
let x = self.g[..(2 * len - 1)].middle_product(&r);
let r = self.g[self.pos + 1 - len..]
.iter()
.cloned()
.rev()
.collect::<Vec<_>>();
let y = self.f[..(2 * len - 1)].middle_product(&r);
if self.pos + x.len() > self.h.len() {
self.h.resize(self.pos + x.len(), T::zero());
}
self.h[self.pos..].add_assign(&x);
self.h[self.pos..].add_assign(&y);
}
self.pos += 1;
self.h[self.pos - 1]
}
// 以降0を仮定した時の添字xの値を求める
// x - pos が小さいかつ同じ添字を頻繁に聞かないことを想定している
pub fn find_assume(&self, x: usize) -> T {
if x < self.pos {
return self.h[x];
}
let mut pos = self.pos;
let mut ans = self.h.get(x).cloned().unwrap_or(T::zero());
while pos <= x {
let a = pos + 2;
let len = 1 << a.trailing_zeros();
if a == len {
for (i, f) in self.f.iter().enumerate() {
if let Some(g) = self.g.get(x - i) {
ans = ans + *f * *g;
}
}
} else {
if x < pos + len {
let f = &self.f;
let g = &self.g;
for i in (pos + 1 - len)..f.len() {
if x - i < f.len() {
ans = ans + f[i] * g[x - i];
ans = ans + g[i] * f[x - i];
}
}
}
}
pos += 1;
}
ans
}
}
// ---------- 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 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 VALID: () = assert!(is_prime(M) && M % 2 == 1 && M < (1 << 30));
const PRIMITIVE_ROOT: u32 = primitive_root(M);
const ORDER: usize = 1 << (M - 1).trailing_zeros();
const fn reduce(x: u64) -> u32 {
let _ = Self::VALID;
let b = (x as u32 * Self::REM) as u64;
let t = x + b * M as u64;
(t >> 32) 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 {
Self(Self::reduce((v % M) 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 {
let mut res = Self::reduce(self.0 as u64);
if res >= M {
res -= M;
}
res
}
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 >= 2 * M {
v -= 2 * 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 += 2 * 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(2 * 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> From<u64> for ModInt<{ M }> {
fn from(val: u64) -> ModInt<{ M }> {
ModInt::new((val % M as u64) as u32)
}
}
impl<const M: u32> From<i64> for ModInt<{ M }> {
fn from(val: i64) -> ModInt<{ M }> {
ModInt::new(val.rem_euclid(M as i64) as u32)
}
}
impl<const M: u32> Into<usize> for ModInt<{ M }> {
fn into(self) -> usize {
self.get() as usize
}
}
// ---------- end modint ----------
// ---------- begin precalc ----------
pub struct Precalc<T> {
fact: Vec<T>,
ifact: Vec<T>,
inv: Vec<T>,
}
impl<T> Precalc<T>
where
T: Copy + Field,
{
pub fn new(size: usize) -> Self {
let mut fact = vec![T::one(); size + 1];
let mut ifact = vec![T::one(); size + 1];
let mut inv = vec![T::one(); size + 1];
let mut mul = T::one();
for i in 2..=size {
mul = mul + T::one();
fact[i] = fact[i - 1] * mul;
}
ifact[size] = T::one() / fact[size];
for i in (2..=size).rev() {
inv[i] = ifact[i] * fact[i - 1];
ifact[i - 1] = ifact[i] * mul;
mul = mul - T::one();
}
Self { fact, ifact, inv }
}
pub fn fact(&self, n: usize) -> T {
self.fact[n]
}
pub fn ifact(&self, n: usize) -> T {
self.ifact[n]
}
pub fn inv(&self, n: usize) -> T {
assert!(0 < n);
self.inv[n]
}
pub fn perm(&self, n: usize, k: usize) -> T {
if k > n {
return T::zero();
}
self.fact[n] * self.ifact[n - k]
}
pub fn binom(&self, n: usize, k: usize) -> T {
if n < k {
return T::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
}
}
pub trait NTT {
fn ntt(&mut self);
fn intt(&mut self);
fn transform(&mut self, len: usize);
fn inverse_transform(&mut self, len: usize);
fn dot_product_ntt(&mut self, rhs: &Self, len: usize);
}
impl<const M: u32> NTT for [ModInt<{ M }>] {
fn ntt(&mut self) {
self.transform(1);
}
fn intt(&mut self) {
self.inverse_transform(1);
}
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 dot_product_ntt(&mut self, rhs: &Self, len: usize) {
let mut buf = [ModInt::zero(); 20];
let buf = &mut buf[..(2 * len - 1)];
let pre = &NTTPrecalcHelper::<{ M }>::A;
let mut now = ModInt::one();
for (i, (f, g)) in self
.chunks_exact_mut(2 * len)
.zip(rhs.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];
}
}
}
// transform でlen=1を指定すればNTTになる
pub trait ArrayConvolution {
type Item;
fn convolution(&self, rhs: &[Self::Item]) -> Vec<Self::Item>;
fn middle_product(&self, a: &[Self::Item]) -> Vec<Self::Item>;
}
pub fn convolution_modulo<const MOD: u32, const A: u32>(
a: &[ModInt<MOD>],
b: &[ModInt<MOD>],
) -> Vec<ModInt<A>> {
let a = a
.iter()
.map(|a| ModInt::<A>::new(a.get()))
.collect::<Vec<_>>();
let b = b
.iter()
.map(|a| ModInt::<A>::new(a.get()))
.collect::<Vec<_>>();
a.convolution(&b)
}
pub fn middle_product_modulo<const MOD: u32, const A: u32>(
a: &[ModInt<MOD>],
b: &[ModInt<MOD>],
) -> Vec<ModInt<A>> {
let a = a
.iter()
.map(|a| ModInt::<A>::new(a.get()))
.collect::<Vec<_>>();
let b = b
.iter()
.map(|a| ModInt::<A>::new(a.get()))
.collect::<Vec<_>>();
a.middle_product(&b)
}
impl<const M: u32> ArrayConvolution for [ModInt<{ M }>] {
type Item = ModInt<{ M }>;
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;
}
if ModInt::<{ M }>::ORDER < k {
const A: u32 = 167772161;
const B: u32 = 469762049;
const C: u32 = 754974721;
assert!(ModInt::<A>::ORDER >= k);
assert!(ModInt::<B>::ORDER >= k);
assert!(ModInt::<C>::ORDER >= k);
const P: u32 = pow_mod(A, B - 2, B);
const Q: u32 = pow_mod(A, C - 2, C);
const R: u32 = pow_mod(B, C - 2, C);
const QR: u32 = (Q as u64 * R as u64 % C as u64) as u32;
const W1: u32 = A;
let w2: u32 = (A as u64 * B as u64 % M as u64) as u32;
let x: Vec<ModInt<A>> = convolution_modulo(self, rhs);
let y: Vec<ModInt<B>> = convolution_modulo(self, rhs);
let z: Vec<ModInt<C>> = convolution_modulo(self, rhs);
let mut ans = vec![ModInt::<{ M }>::zero(); x.len()];
for (((ans, x), y), z) in ans.iter_mut().zip(x).zip(y).zip(z) {
let a = x.get();
let b = ((y.get() + B - a) as u64 * P as u64 % B as u64) as u32;
let c = (((z.get() + C - a) as u64 * QR as u64 + (C - b) as u64 * R as u64)
% C as u64) as u32;
*ans = (a as u64 + b as u64 * W1 as u64 + c as u64 * w2 as u64).into();
}
return ans;
}
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);
f.dot_product_ntt(&g, len);
f.inverse_transform(len);
f.truncate(self.len() + rhs.len() - 1);
f
}
fn middle_product(&self, rhs: &[Self::Item]) -> Vec<Self::Item> {
assert!(self.len() >= rhs.len());
if self.len() - rhs.len() <= 32 {
return self
.windows(rhs.len())
.map(|a| {
a.iter()
.zip(rhs.iter())
.fold(ModInt::zero(), |s, p| s + *p.0 * *p.1)
})
.collect();
}
const PARAM: usize = 10;
let size = self.len();
let mut k = 0;
while (size + (1 << k) - 1) >> k > PARAM {
k += 1;
}
if ModInt::<{ M }>::ORDER < k {
const A: u32 = 167772161;
const B: u32 = 469762049;
const C: u32 = 754974721;
assert!(ModInt::<A>::ORDER >= k);
assert!(ModInt::<B>::ORDER >= k);
assert!(ModInt::<C>::ORDER >= k);
const P: u32 = pow_mod(A, B - 2, B);
const Q: u32 = pow_mod(A, C - 2, C);
const R: u32 = pow_mod(B, C - 2, C);
const QR: u32 = (Q as u64 * R as u64 % C as u64) as u32;
const W1: u32 = A;
let w2: u32 = (A as u64 * B as u64 % M as u64) as u32;
let x: Vec<ModInt<A>> = middle_product_modulo(self, rhs);
let y: Vec<ModInt<B>> = middle_product_modulo(self, rhs);
let z: Vec<ModInt<C>> = middle_product_modulo(self, rhs);
let mut ans = vec![ModInt::<{ M }>::zero(); x.len()];
for (((ans, x), y), z) in ans.iter_mut().zip(x).zip(y).zip(z) {
let a = x.get();
let b = ((y.get() + B - a) as u64 * P as u64 % B as u64) as u32;
let c = (((z.get() + C - a) as u64 * QR as u64 + (C - b) as u64 * R as u64)
% C as u64) as u32;
*ans = (a as u64 + b as u64 * W1 as u64 + c as u64 * w2 as u64).into();
}
return ans;
}
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);
g[..rhs.len()].reverse();
f.transform(len);
g.transform(len);
f.dot_product_ntt(&g, len);
f.inverse_transform(len);
(rhs.len()..=self.len()).map(|i| f[i - 1]).collect()
}
}
pub trait PolynomialOperation {
type Item;
fn eval(&self, x: Self::Item) -> Self::Item;
fn derivative(&self) -> Vec<Self::Item>;
fn integral(&self) -> Vec<Self::Item>;
}
impl<T> PolynomialOperation for [T]
where
T: Field + Copy,
{
type Item = T;
fn eval(&self, x: Self::Item) -> Self::Item {
self.iter().rfold(T::zero(), |s, a| s * x + *a)
}
fn derivative(&self) -> Vec<Self::Item> {
if self.len() <= 1 {
return vec![];
}
self[1..]
.iter()
.scan(T::one(), |s, a| {
let res = *a * *s;
*s = *s + T::one();
Some(res)
})
.collect()
}
fn integral(&self) -> Vec<Self::Item> {
if self.is_empty() {
return vec![];
}
let mut inv = vec![T::one(); self.len() + 1];
let mut val = T::zero();
for i in 1..inv.len() {
val = val + T::one();
inv[i] = val * inv[i - 1];
}
let mut iprod = T::one() / inv[self.len()];
for i in (1..inv.len()).rev() {
inv[i] = iprod * inv[i - 1] * self[i - 1];
iprod = iprod * val;
val = val - T::one();
}
inv[0] = T::zero();
inv
}
}
pub trait FPSOperation {
type Item;
fn inverse(&self, n: usize) -> Vec<Self::Item>;
fn log(&self, n: usize) -> Vec<Self::Item>;
fn exp(&self, n: usize) -> Vec<Self::Item>;
}
impl<T> FPSOperation for [T]
where
T: Field + Copy,
[T]: ArrayConvolution<Item = T>,
{
type Item = T;
fn inverse(&self, n: usize) -> Vec<Self::Item> {
if n == 0 {
return vec![];
}
assert!(self.len() > 0 && !self[0].is_zero());
let mut g = Vec::with_capacity(n);
g.push(T::one() / self[0]);
while g.len() < n {
let size = g.len();
let up = (2 * size).min(n);
let gg = g.convolution(&g);
let mut h = gg.convolution(&self[..up.min(self.len())]);
h.resize(up, T::zero());
g.extend(h[size..up].iter().map(|v| -*v));
}
g
}
fn log(&self, n: usize) -> Vec<Self::Item> {
assert!(self.len() > 0 && self[0].is_one());
if n == 0 {
return vec![];
}
let mut res = self.derivative().convolution(&self.inverse(n));
res.truncate(n - 1);
res.integral()
}
fn exp(&self, n: usize) -> Vec<Self::Item> {
if n == 0 {
return vec![];
}
if self.is_empty() {
let mut res = vec![T::zero(); n];
res[0] = T::one();
return res;
}
assert!(self.len() > 0 && self[0].is_zero());
let mut g = Vec::with_capacity(n);
g.push(T::one());
while g.len() < n {
let size = g.len();
let up = (2 * size).min(n);
let lg = g.log(up);
let rhs = self[..up.min(self.len())].sub(&lg);
let mut h = g.convolution(&rhs);
h.resize(up, T::zero());
g.extend(h[size..up].iter().cloned());
}
g
}
}
// ---------- end array op ----------
// ---------- begin trait ----------
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 Group: Zero + Sub<Output = Self> + Neg<Output = Self> {}
pub trait SemiRing: Zero + One {}
pub trait Ring: SemiRing + Group {}
pub trait Field: Ring + Div<Output = Self> {}
impl<T> Group for T where T: Zero + Sub<Output = Self> + Neg<Output = Self> {}
impl<T> SemiRing for T where T: Zero + One {}
impl<T> Ring for T where T: SemiRing + Group {}
impl<T> Field for T where T: Ring + Div<Output = Self> {}
pub fn zero<T: Zero>() -> T {
T::zero()
}
pub fn one<T: One>() -> T {
T::one()
}
pub fn pow<T: One + Clone>(mut r: T, mut n: usize) -> T {
let mut t = one();
while n > 0 {
if n & 1 == 1 {
t = t * r.clone();
}
r = r.clone() * r;
n >>= 1;
}
t
}
pub fn pow_sum<T: SemiRing + Clone>(r: T, n: usize) -> T {
if n == 0 {
T::zero()
} else if n & 1 == 1 {
T::one() + r.clone() * pow_sum(r, n - 1)
} else {
let a = T::one() + r.clone();
let b = r.clone() * r;
a * pow_sum(b, n / 2)
}
}
// ---------- end trait ----------
// ---------- taylor shift ----------
// f(x) とcを受け取って f(x+c) を返す
pub trait TaylorShift {
type Item;
fn taylor_shift(&self, c: Self::Item) -> Vec<Self::Item>;
}
impl<T> TaylorShift for [T]
where
T: Copy + Field,
[T]: ArrayConvolution<Item = T>,
{
type Item = T;
fn taylor_shift(&self, c: Self::Item) -> Vec<Self::Item> {
if self.is_empty() || c.is_zero() {
return Vec::from(self);
}
let mut fact = vec![T::one(); self.len()];
let mut val = T::zero();
for i in 1..fact.len() {
val = val + T::one();
fact[i] = fact[i - 1] * val;
}
let mut ifact = vec![T::one(); self.len()];
ifact[self.len() - 1] = T::one() / fact[self.len() - 1];
for i in (1..fact.len()).rev() {
ifact[i - 1] = ifact[i] * val;
val = val - T::one();
}
let mut a = Vec::from(self);
for (a, f) in a.iter_mut().zip(fact.iter()) {
*a = *a * *f;
}
a.reverse();
let mut pow = T::one();
for (f, i) in fact.iter_mut().zip(ifact.iter()) {
*f = *i * pow;
pow = pow * c;
}
a = a.convolution(&fact);
a.truncate(self.len());
a.reverse();
for (a, i) in a.iter_mut().zip(ifact.iter()) {
*a = *a * *i;
}
a
}
}
// ---------- taylor shift ----------
pub fn composition_inverse<T>(mut f: Vec<T>, n: usize) -> Vec<T>
where
T: Field + From<usize> + Copy + std::fmt::Debug,
[T]: ArrayConvolution<Item = T>,
{
assert!(f.len() >= 2 && f[0].is_zero() && !f[1].is_zero());
let f1inv = T::one() / f[1];
if n <= 2 {
let mut res = vec![T::zero(), f1inv];
res.truncate(n);
return res;
}
f.truncate(n);
let n = n - 1;
for f in f.iter_mut() {
*f = *f * f1inv;
}
let mut de = Poly2d::new(vec![T::zero(); 2 * (n + 1)], 2, n + 1);
de[0][0] = T::one();
for (i, f) in f.iter().enumerate() {
de[1][i] = -*f;
}
let mut nu = Poly2d::new(vec![T::one()], 1, 1);
for j in 0.. {
let mut p = de.clone();
for i in 0..p.h {
for j in (1..p.w).step_by(2) {
p[i][j] = -p[i][j];
}
}
nu = nu.conv(&p);
nu = nu
.resize(nu.h, nu.w.min((n >> j) + 1))
.clip((0, 1), (n >> j & 1, 2));
if n >> (j + 1) == 0 {
break;
}
let mut s = p;
let mut w = de.w;
if w % 2 == 0 {
de = de.resize(de.h, w - 1);
s = s.resize(de.h, w - 1);
w -= 1;
}
let a = de.a.convolution(&s.a);
de = Poly2d::new(a, de.h + s.h - 1, w).clip((0, 1), (0, 2));
}
let pow = (0..=n).map(|i| nu[i][0]).collect::<Vec<_>>();
let mut g = vec![T::zero(); n];
for i in 1..=n {
g[n - i] = T::from(n) * pow[i] / T::from(i);
}
g = g.log(n);
let v = -T::one() / T::from(n);
for g in g.iter_mut() {
*g = *g * v;
}
g = g.exp(n);
g.insert(0, T::zero());
let mut pow = T::one();
for g in g.iter_mut() {
*g = *g * pow;
pow = pow * f1inv;
}
g
}
// f(g(x)) = sum_i f_i (g_0 + g(x))^i
// sum_i sum_{0 <= j <= i} f_i C(i, j) g_0^j g(x)^j
// f(g(x)) の [0, n) 次を求める
pub fn composition_of_fps<T>(mut f: Vec<T>, mut g: Vec<T>, n: usize) -> Vec<T>
where
T: Field + Copy + std::fmt::Debug,
[T]: ArrayConvolution<Item = T>,
{
if f.is_empty() || n == 0 {
return vec![T::zero(); n];
}
if g.len() > 0 && !g[0].is_zero() {
f = f.taylor_shift(std::mem::replace(&mut g[0], T::zero()));
}
if g.iter().position(|g| !g.is_zero()).map_or(true, |x| x >= n) {
let mut res = vec![T::zero(); n];
res[0] = f[0];
return res;
}
let mut memo = vec![];
let mut de = Poly2d::new(vec![T::zero(); 2 * n], 2, n);
de[0][0] = T::one();
for (i, g) in g.iter().enumerate().take(n) {
de[1][i] = -*g;
}
let mut deg = 1;
while deg < n {
let mut s = de.clone();
for s in s.a.chunks_exact_mut(de.w) {
for s in s[1..].iter_mut().step_by(2) {
*s = -*s;
}
}
memo.push(s.clone());
if 2 * deg >= n {
break;
}
let w = de.w;
if w % 2 == 0 {
de = de.resize(de.h, w - 1);
s = s.resize(de.h, w - 1);
}
let a = de.a.convolution(&s.a);
de = Poly2d::new(a, de.h + de.h - 1, de.w);
de = de
.resize(de.h, ((n + deg - 1) / deg).min(de.w))
.clip((0, 1), (0, 2));
deg <<= 1;
}
f.resize(n, T::zero());
let h = f.len();
f.reverse();
let mut f = Poly2d::new(f, h, 1);
while let Some(mut m) = memo.pop() {
let mut nf = Poly2d::new(vec![T::zero(); f.h * (2 * f.w - 1)], f.h, 2 * f.w - 1);
for i in 0..f.h {
for j in 0..f.w {
nf[i][2 * j] = f[i][j];
}
}
if false || f.h != 2 * deg {
f = m.conv(&nf);
let ylow = (nf.h - 1).saturating_sub(deg - 1);
let xup = f.w.min((n + deg - 1) / deg);
f = f.resize(nf.h, xup).clip((ylow, 1), (0, 1));
} else {
let fw = nf.w;
let mw = m.w;
let w = m.w + nf.w - 1;
nf = nf.resize(nf.h, w);
nf.a.rotate_right(mw - 1);
nf.a.extend((1..mw).map(|_| T::zero()));
m = m.resize(m.h, w);
m.a.reverse();
m.a.rotate_left(fw - 1);
for _ in 1..fw {
m.a.pop();
}
nf.a.reverse();
m.a.reverse();
let mut a = nf.a.middle_product(&m.a);
a.reverse();
let nw = ((n + deg - 1) / deg).min(w);
let a = a
.chunks(w)
.flat_map(|a| a.iter().cloned().take(nw))
.collect::<Vec<_>>();
f = Poly2d::new(a, deg, nw);
}
deg >>= 1;
}
f.a.truncate(n);
f.a.resize(n, T::zero());
f.a
}
#[derive(Clone, Debug)]
struct Poly2d<T> {
a: Vec<T>,
h: usize,
w: usize,
}
impl<T> Poly2d<T>
where
T: Copy + Ring,
[T]: ArrayConvolution<Item = T>,
{
pub fn zero() -> Self {
Self::new(vec![], 0, 0)
}
pub fn new(a: Vec<T>, h: usize, w: usize) -> Self {
let mut res = vec![T::zero(); h * w];
for (res, a) in res.chunks_exact_mut(w).zip(a.chunks(w)) {
let l = w.min(a.len());
res[..l].copy_from_slice(&a[..l]);
}
Self { a: res, h, w }
}
pub fn resize(&self, h: usize, w: usize) -> Self {
if h * w == 0 {
return Self::new(vec![], 0, 0);
}
let mut a = vec![T::zero(); h * w];
let l = self.w.min(w);
for (a, b) in a.chunks_exact_mut(w).zip(self.a.chunks(self.w)) {
a[..l].copy_from_slice(&b[..l]);
}
Self::new(a, h, w)
}
fn conv(&self, rhs: &Self) -> Self {
if self.is_empty() {
return rhs.clone();
}
if rhs.is_empty() {
return rhs.clone();
}
let nw = self.w + rhs.w - 1;
let mut a = self.resize(self.h, nw);
let mut b = rhs.resize(rhs.h, nw);
for _ in 1..rhs.w {
a.a.pop();
}
for _ in 1..self.w {
b.a.pop();
}
Self::new(a.a.convolution(&b.a), self.h + rhs.h - 1, nw)
}
fn clip(&self, row: (usize, usize), col: (usize, usize)) -> Self {
if row.0 >= self.h || col.0 >= self.w {
return Self::zero();
}
let h = (self.h - row.0 + row.1 - 1) / row.1;
let w = (self.w - col.0 + col.1 - 1) / col.1;
let mut res = Self::new(vec![T::zero(); h * w], h, w);
for i in 0..h {
for j in 0..w {
res[i][j] = self[row.0 + row.1 * i][col.0 + col.1 * j];
}
}
res
}
fn is_empty(&self) -> bool {
self.a.is_empty()
}
}
impl<T> Index<usize> for Poly2d<T> {
type Output = [T];
fn index(&self, x: usize) -> &Self::Output {
assert!(x < self.h);
let l = x * self.w;
let r = l + self.w;
&self.a[l..r]
}
}
impl<T> IndexMut<usize> for Poly2d<T> {
fn index_mut(&mut self, x: usize) -> &mut Self::Output {
assert!(x < self.h);
let l = x * self.w;
let r = l + self.w;
&mut self.a[l..r]
}
}