結果
| 問題 | No.3559 +A,-B |
| コンテスト | |
| ユーザー |
 akkey
|
| 提出日時 | 2026-05-29 23:44:24 |
| 言語 | Rust (1.94.0 + proconio + num + itertools) |
| 結果 |
TLE
|
| 実行時間 | - |
| コード長 | 61,037 bytes |
| 記録 | |
| コンパイル時間 | 1,006 ms |
| コンパイル使用メモリ | 228,080 KB |
| 実行使用メモリ | 12,416 KB |
| 最終ジャッジ日時 | 2026-05-29 23:45:14 |
| 合計ジャッジ時間 | 6,478 ms |
|
ジャッジサーバーID （参考情報） |
judge2_0 / judge1_0 |
（要ログイン）
| サブタスク | 配点 | 結果 |
|---|---|---|
| 部分点1 | 10 % | AC * 4 |
| 部分点2 | 60 % | AC * 2 TLE * 1 -- * 4 |
| 部分点3 | 30 % | AC * 6 TLE * 1 -- * 13 |
| 合計 | 10 点 |
ソースコード
#![allow(non_snake_case, unused_imports, unused_macros)]
use itertools::Itertools;
use proconio::{fastout, input, marker::Usize1};
macro_rules! debug {
($($a:expr),* $(,)*) => {
#[cfg(debug_assertions)]
eprintln!(concat!($("| ", stringify!($a), "={:?} "),*, "|"), $(&$a),*);
};
}
macro_rules! ndvec {
($v:expr; $n:expr) => {
vec![$v; $n]
};
($v:expr; $n:expr, $($ns:expr),+) => {
vec![ndvec![$v; $($ns),+]; $n]
};
}
macro_rules! yes_no {
($e:expr) => {
if $e {
println!("Yes");
} else {
println!("No");
}
};
}
#[fastout]
fn main() {
input!{
t: usize,
nxyab: [(usize, isize, isize, isize, isize); t]
}
for (n, mut x, mut y, a, b) in nxyab {
for _ in 0..n {
if x + y >= 0 {
y -= b;
} else {
x += a;
}
}
println!("{x} {y}");
}
}
#[allow(unused)]
fn partition_point<I: num::Integer + Copy>(mut l: I, mut r: I, mut f: impl FnMut(I) -> bool) {
let one = I::one();
let two = one + one;
while l < r {
let p = (r - l) / two + l;
if f(p) {
l = p + one;
} else {
r = p;
}
}
}
#[allow(dead_code)]
pub(crate) fn ceil_pow2(n: u32) -> u32 {
32 - n.saturating_sub(1).leading_zeros()
}
use std::{
fmt,
iter::{Product, Sum},
ops::{
Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div,
DivAssign, Mul, MulAssign, Not, Rem, RemAssign, Shl, ShlAssign, Shr, ShrAssign, Sub,
SubAssign,
},
};
pub trait Integral:
'static
+ Send
+ Sync
+ Copy
+ Ord
+ Not<Output = Self>
+ Add<Output = Self>
+ Sub<Output = Self>
+ Mul<Output = Self>
+ Div<Output = Self>
+ Rem<Output = Self>
+ AddAssign
+ SubAssign
+ MulAssign
+ DivAssign
+ RemAssign
+ Sum
+ Product
+ BitOr<Output = Self>
+ BitAnd<Output = Self>
+ BitXor<Output = Self>
+ BitOrAssign
+ BitAndAssign
+ BitXorAssign
+ Shl<Output = Self>
+ Shr<Output = Self>
+ ShlAssign
+ ShrAssign
+ fmt::Display
+ fmt::Debug
+ fmt::Binary
+ fmt::Octal
+ Zero
+ One
+ BoundedBelow
+ BoundedAbove
{
}
/// Class that has additive identity element
pub trait Zero {
/// The additive identity element
fn zero() -> Self;
}
/// Class that has multiplicative identity element
pub trait One {
/// The multiplicative identity element
fn one() -> Self;
}
pub trait BoundedBelow {
fn min_value() -> Self;
}
pub trait BoundedAbove {
fn max_value() -> Self;
}
macro_rules! impl_integral {
($($ty:ty),*) => {
$(
impl Zero for $ty {
#[inline]
fn zero() -> Self {
0
}
}
impl One for $ty {
#[inline]
fn one() -> Self {
1
}
}
impl BoundedBelow for $ty {
#[inline]
fn min_value() -> Self {
Self::MIN
}
}
impl BoundedAbove for $ty {
#[inline]
fn max_value() -> Self {
Self::MAX
}
}
impl Integral for $ty {}
)*
};
}
impl_integral!(
i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize
);
#[allow(dead_code)]
#[derive(Default)]
pub(crate) struct SimpleQueue<T> {
payload: Vec<T>,
pos: usize,
}
#[allow(dead_code)]
impl<T> SimpleQueue<T> {
pub(crate) fn reserve(&mut self, n: usize) {
if n > self.payload.len() {
self.payload.reserve(n - self.payload.len());
}
}
pub(crate) fn size(&self) -> usize {
self.payload.len() - self.pos
}
pub(crate) fn empty(&self) -> bool {
self.pos == self.payload.len()
}
pub(crate) fn push(&mut self, t: T) {
self.payload.push(t);
}
// Do we need mutable version?
pub(crate) fn front(&self) -> Option<&T> {
if self.pos < self.payload.len() {
Some(&self.payload[self.pos])
} else {
None
}
}
pub(crate) fn clear(&mut self) {
self.payload.clear();
self.pos = 0;
}
pub(crate) fn pop(&mut self) -> Option<&T> {
if self.pos < self.payload.len() {
self.pos += 1;
Some(&self.payload[self.pos - 1])
} else {
None
}
}
}
use std::cmp::{max, min};
use std::convert::Infallible;
use std::iter::FromIterator;
use std::marker::PhantomData;
use std::ops::{Bound, RangeBounds};
// TODO Should I split monoid-related traits to another module?
pub trait Monoid {
type S: Clone;
fn identity() -> Self::S;
fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S;
}
pub struct Max<S>(Infallible, PhantomData<fn() -> S>);
impl<S> Monoid for Max<S>
where
S: Copy + Ord + BoundedBelow,
{
type S = S;
fn identity() -> Self::S {
S::min_value()
}
fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
max(*a, *b)
}
}
pub struct Min<S>(Infallible, PhantomData<fn() -> S>);
impl<S> Monoid for Min<S>
where
S: Copy + Ord + BoundedAbove,
{
type S = S;
fn identity() -> Self::S {
S::max_value()
}
fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
min(*a, *b)
}
}
pub struct Additive<S>(Infallible, PhantomData<fn() -> S>);
impl<S> Monoid for Additive<S>
where
S: Copy + Add<Output = S> + Zero,
{
type S = S;
fn identity() -> Self::S {
S::zero()
}
fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
*a + *b
}
}
pub struct Multiplicative<S>(Infallible, PhantomData<fn() -> S>);
impl<S> Monoid for Multiplicative<S>
where
S: Copy + Mul<Output = S> + One,
{
type S = S;
fn identity() -> Self::S {
S::one()
}
fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
*a * *b
}
}
pub struct BitwiseOr<S>(Infallible, PhantomData<fn() -> S>);
impl<S> Monoid for BitwiseOr<S>
where
S: Copy + BitOr<Output = S> + Zero,
{
type S = S;
fn identity() -> Self::S {
S::zero()
}
fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
*a | *b
}
}
pub struct BitwiseAnd<S>(Infallible, PhantomData<fn() -> S>);
impl<S> Monoid for BitwiseAnd<S>
where
S: Copy + BitAnd<Output = S> + Not<Output = S> + Zero,
{
type S = S;
fn identity() -> Self::S {
!S::zero()
}
fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
*a & *b
}
}
pub struct BitwiseXor<S>(Infallible, PhantomData<fn() -> S>);
impl<S> Monoid for BitwiseXor<S>
where
S: Copy + BitXor<Output = S> + Zero,
{
type S = S;
fn identity() -> Self::S {
S::zero()
}
fn binary_operation(a: &Self::S, b: &Self::S) -> Self::S {
*a ^ *b
}
}
impl<M: Monoid> Default for Segtree<M> {
fn default() -> Self {
Segtree::new(0)
}
}
impl<M: Monoid> Segtree<M> {
pub fn new(n: usize) -> Segtree<M> {
vec![M::identity(); n].into()
}
}
impl<M: Monoid> From<Vec<M::S>> for Segtree<M> {
fn from(v: Vec<M::S>) -> Self {
let n = v.len();
let log = ceil_pow2(n as u32) as usize;
let size = 1 << log;
let mut d = vec![M::identity(); 2 * size];
d[size..][..n].clone_from_slice(&v);
let mut ret = Segtree { n, size, log, d };
for i in (1..size).rev() {
ret.update(i);
}
ret
}
}
impl<M: Monoid> FromIterator<M::S> for Segtree<M> {
fn from_iter<T: IntoIterator<Item = M::S>>(iter: T) -> Self {
let v = iter.into_iter().collect::<Vec<_>>();
v.into()
}
}
impl<M: Monoid> Segtree<M> {
pub fn set(&mut self, mut p: usize, x: M::S) {
assert!(p < self.n);
p += self.size;
self.d[p] = x;
for i in 1..=self.log {
self.update(p >> i);
}
}
pub fn get(&self, p: usize) -> M::S {
assert!(p < self.n);
self.d[p + self.size].clone()
}
pub fn get_slice(&self) -> &[M::S] {
&self.d[self.size..][..self.n]
}
pub fn prod<R>(&self, range: R) -> M::S
where
R: RangeBounds<usize>,
{
// Trivial optimization
if range.start_bound() == Bound::Unbounded && range.end_bound() == Bound::Unbounded {
return self.all_prod();
}
let mut r = match range.end_bound() {
Bound::Included(r) => r + 1,
Bound::Excluded(r) => *r,
Bound::Unbounded => self.n,
};
let mut l = match range.start_bound() {
Bound::Included(l) => *l,
Bound::Excluded(l) => l + 1,
// TODO: There are another way of optimizing [0..r)
Bound::Unbounded => 0,
};
assert!(l <= r && r <= self.n);
let mut sml = M::identity();
let mut smr = M::identity();
l += self.size;
r += self.size;
while l < r {
if l & 1 != 0 {
sml = M::binary_operation(&sml, &self.d[l]);
l += 1;
}
if r & 1 != 0 {
r -= 1;
smr = M::binary_operation(&self.d[r], &smr);
}
l >>= 1;
r >>= 1;
}
M::binary_operation(&sml, &smr)
}
pub fn all_prod(&self) -> M::S {
self.d[1].clone()
}
pub fn max_right<F>(&self, mut l: usize, f: F) -> usize
where
F: Fn(&M::S) -> bool,
{
assert!(l <= self.n);
assert!(f(&M::identity()));
if l == self.n {
return self.n;
}
l += self.size;
let mut sm = M::identity();
while {
// do
while l.is_multiple_of(2) {
l >>= 1;
}
if !f(&M::binary_operation(&sm, &self.d[l])) {
while l < self.size {
l *= 2;
let res = M::binary_operation(&sm, &self.d[l]);
if f(&res) {
sm = res;
l += 1;
}
}
return l - self.size;
}
sm = M::binary_operation(&sm, &self.d[l]);
l += 1;
// while
{
let l = l as isize;
(l & -l) != l
}
} {}
self.n
}
pub fn min_left<F>(&self, mut r: usize, f: F) -> usize
where
F: Fn(&M::S) -> bool,
{
assert!(r <= self.n);
assert!(f(&M::identity()));
if r == 0 {
return 0;
}
r += self.size;
let mut sm = M::identity();
while {
// do
r -= 1;
while r > 1 && r % 2 == 1 {
r >>= 1;
}
if !f(&M::binary_operation(&self.d[r], &sm)) {
while r < self.size {
r = 2 * r + 1;
let res = M::binary_operation(&self.d[r], &sm);
if f(&res) {
sm = res;
r -= 1;
}
}
return r + 1 - self.size;
}
sm = M::binary_operation(&self.d[r], &sm);
// while
{
let r = r as isize;
(r & -r) != r
}
} {}
0
}
fn update(&mut self, k: usize) {
self.d[k] = M::binary_operation(&self.d[2 * k], &self.d[2 * k + 1]);
}
}
// Maybe we can use this someday
// ```
// for i in 0..=self.log {
// for j in 0..1 << i {
// print!("{}\t", self.d[(1 << i) + j]);
// }
// println!();
// }
// ```
#[derive(Clone)]
pub struct Segtree<M>
where
M: Monoid,
{
// variable name is _n in original library
n: usize,
size: usize,
log: usize,
d: Vec<M::S>,
}
pub trait MapMonoid {
type M: Monoid;
type F: Clone;
// type S = <Self::M as Monoid>::S;
fn identity_element() -> <Self::M as Monoid>::S {
Self::M::identity()
}
fn binary_operation(
a: &<Self::M as Monoid>::S,
b: &<Self::M as Monoid>::S,
) -> <Self::M as Monoid>::S {
Self::M::binary_operation(a, b)
}
fn identity_map() -> Self::F;
fn mapping(f: &Self::F, x: &<Self::M as Monoid>::S) -> <Self::M as Monoid>::S;
fn composition(f: &Self::F, g: &Self::F) -> Self::F;
}
impl<F: MapMonoid> Default for LazySegtree<F> {
fn default() -> Self {
Self::new(0)
}
}
impl<F: MapMonoid> LazySegtree<F> {
pub fn new(n: usize) -> Self {
vec![F::identity_element(); n].into()
}
}
impl<F: MapMonoid> From<Vec<<F::M as Monoid>::S>> for LazySegtree<F> {
fn from(v: Vec<<F::M as Monoid>::S>) -> Self {
let n = v.len();
let log = ceil_pow2(n as u32) as usize;
let size = 1 << log;
let mut d = vec![F::identity_element(); 2 * size];
let lz = vec![F::identity_map(); size];
d[size..(size + n)].clone_from_slice(&v);
let mut ret = LazySegtree {
n,
size,
log,
d,
lz,
};
for i in (1..size).rev() {
ret.update(i);
}
ret
}
}
impl<F: MapMonoid> LazySegtree<F> {
pub fn set(&mut self, mut p: usize, x: <F::M as Monoid>::S) {
assert!(p < self.n);
p += self.size;
for i in (1..=self.log).rev() {
self.push(p >> i);
}
self.d[p] = x;
for i in 1..=self.log {
self.update(p >> i);
}
}
pub fn get(&mut self, mut p: usize) -> <F::M as Monoid>::S {
assert!(p < self.n);
p += self.size;
for i in (1..=self.log).rev() {
self.push(p >> i);
}
self.d[p].clone()
}
pub fn prod<R>(&mut self, range: R) -> <F::M as Monoid>::S
where
R: RangeBounds<usize>,
{
// Trivial optimization
if range.start_bound() == Bound::Unbounded && range.end_bound() == Bound::Unbounded {
return self.all_prod();
}
let mut r = match range.end_bound() {
Bound::Included(r) => r + 1,
Bound::Excluded(r) => *r,
Bound::Unbounded => self.n,
};
let mut l = match range.start_bound() {
Bound::Included(l) => *l,
Bound::Excluded(l) => l + 1,
// TODO: There are another way of optimizing [0..r)
Bound::Unbounded => 0,
};
assert!(l <= r && r <= self.n);
if l == r {
return F::identity_element();
}
l += self.size;
r += self.size;
for i in (1..=self.log).rev() {
if ((l >> i) << i) != l {
self.push(l >> i);
}
if ((r >> i) << i) != r {
self.push(r >> i);
}
}
let mut sml = F::identity_element();
let mut smr = F::identity_element();
while l < r {
if l & 1 != 0 {
sml = F::binary_operation(&sml, &self.d[l]);
l += 1;
}
if r & 1 != 0 {
r -= 1;
smr = F::binary_operation(&self.d[r], &smr);
}
l >>= 1;
r >>= 1;
}
F::binary_operation(&sml, &smr)
}
pub fn all_prod(&self) -> <F::M as Monoid>::S {
self.d[1].clone()
}
pub fn apply(&mut self, mut p: usize, f: F::F) {
assert!(p < self.n);
p += self.size;
for i in (1..=self.log).rev() {
self.push(p >> i);
}
self.d[p] = F::mapping(&f, &self.d[p]);
for i in 1..=self.log {
self.update(p >> i);
}
}
pub fn apply_range<R>(&mut self, range: R, f: F::F)
where
R: RangeBounds<usize>,
{
let mut r = match range.end_bound() {
Bound::Included(r) => r + 1,
Bound::Excluded(r) => *r,
Bound::Unbounded => self.n,
};
let mut l = match range.start_bound() {
Bound::Included(l) => *l,
Bound::Excluded(l) => l + 1,
// TODO: There are another way of optimizing [0..r)
Bound::Unbounded => 0,
};
assert!(l <= r && r <= self.n);
if l == r {
return;
}
l += self.size;
r += self.size;
for i in (1..=self.log).rev() {
if ((l >> i) << i) != l {
self.push(l >> i);
}
if ((r >> i) << i) != r {
self.push((r - 1) >> i);
}
}
{
let l2 = l;
let r2 = r;
while l < r {
if l & 1 != 0 {
self.all_apply(l, f.clone());
l += 1;
}
if r & 1 != 0 {
r -= 1;
self.all_apply(r, f.clone());
}
l >>= 1;
r >>= 1;
}
l = l2;
r = r2;
}
for i in 1..=self.log {
if ((l >> i) << i) != l {
self.update(l >> i);
}
if ((r >> i) << i) != r {
self.update((r - 1) >> i);
}
}
}
pub fn max_right<G>(&mut self, mut l: usize, g: G) -> usize
where
G: Fn(<F::M as Monoid>::S) -> bool,
{
assert!(l <= self.n);
assert!(g(F::identity_element()));
if l == self.n {
return self.n;
}
l += self.size;
for i in (1..=self.log).rev() {
self.push(l >> i);
}
let mut sm = F::identity_element();
while {
// do
while l.is_multiple_of(2) {
l >>= 1;
}
if !g(F::binary_operation(&sm, &self.d[l])) {
while l < self.size {
self.push(l);
l *= 2;
let res = F::binary_operation(&sm, &self.d[l]);
if g(res.clone()) {
sm = res;
l += 1;
}
}
return l - self.size;
}
sm = F::binary_operation(&sm, &self.d[l]);
l += 1;
//while
{
let l = l as isize;
(l & -l) != l
}
} {}
self.n
}
pub fn min_left<G>(&mut self, mut r: usize, g: G) -> usize
where
G: Fn(<F::M as Monoid>::S) -> bool,
{
assert!(r <= self.n);
assert!(g(F::identity_element()));
if r == 0 {
return 0;
}
r += self.size;
for i in (1..=self.log).rev() {
self.push((r - 1) >> i);
}
let mut sm = F::identity_element();
while {
// do
r -= 1;
while r > 1 && !r.is_multiple_of(2) {
r >>= 1;
}
if !g(F::binary_operation(&self.d[r], &sm)) {
while r < self.size {
self.push(r);
r = 2 * r + 1;
let res = F::binary_operation(&self.d[r], &sm);
if g(res.clone()) {
sm = res;
r -= 1;
}
}
return r + 1 - self.size;
}
sm = F::binary_operation(&self.d[r], &sm);
// while
{
let r = r as isize;
(r & -r) != r
}
} {}
0
}
}
#[derive(Clone)]
pub struct LazySegtree<F>
where
F: MapMonoid,
{
n: usize,
size: usize,
log: usize,
d: Vec<<F::M as Monoid>::S>,
lz: Vec<F::F>,
}
impl<F> LazySegtree<F>
where
F: MapMonoid,
{
fn update(&mut self, k: usize) {
self.d[k] = F::binary_operation(&self.d[2 * k], &self.d[2 * k + 1]);
}
fn all_apply(&mut self, k: usize, f: F::F) {
self.d[k] = F::mapping(&f, &self.d[k]);
if k < self.size {
self.lz[k] = F::composition(&f, &self.lz[k]);
}
}
fn push(&mut self, k: usize) {
self.all_apply(2 * k, self.lz[k].clone());
self.all_apply(2 * k + 1, self.lz[k].clone());
self.lz[k] = F::identity_map();
}
}
// TODO is it useful?
use std::fmt::{Debug, Error, Formatter, Write};
impl<F> Debug for LazySegtree<F>
where
F: MapMonoid,
F::F: Debug,
<F::M as Monoid>::S: Debug,
{
fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
for i in 0..self.log {
for j in 0..1 << i {
f.write_fmt(format_args!(
"{:?}[{:?}]\t",
self.d[(1 << i) + j],
self.lz[(1 << i) + j]
))?;
}
f.write_char('\n')?;
}
for i in 0..self.size {
f.write_fmt(format_args!("{:?}\t", self.d[self.size + i]))?;
}
Ok(())
}
}
mod internal_math {
// remove this after dependencies has been added
#![allow(dead_code)]
use std::{mem::swap, num::Wrapping as W};
/// # Arguments
/// * `m` `1 <= m`
///
/// # Returns
/// x mod m
/* const */
pub(crate) fn safe_mod(mut x: i64, m: i64) -> i64 {
x %= m;
if x < 0 {
x += m;
}
x
}
/// Fast modular by barrett reduction
/// Reference: https://en.wikipedia.org/wiki/Barrett_reduction
/// NOTE: reconsider after Ice Lake
pub(crate) struct Barrett {
pub(crate) _m: u32,
pub(crate) im: u64,
}
impl Barrett {
/// # Arguments
/// * `m` `1 <= m`
/// (Note: `m <= 2^31` should also hold, which is undocumented in the original library.
/// See the [pull reqeust commment](https://github.com/rust-lang-ja/ac-library-rs/pull/3#discussion_r484661007)
/// for more details.)
pub(crate) fn new(m: u32) -> Barrett {
Barrett {
_m: m,
im: (-1i64 as u64 / m as u64).wrapping_add(1),
}
}
/// # Returns
/// `m`
pub(crate) fn umod(&self) -> u32 {
self._m
}
/// # Parameters
/// * `a` `0 <= a < m`
/// * `b` `0 <= b < m`
///
/// # Returns
/// a * b % m
#[allow(clippy::many_single_char_names)]
pub(crate) fn mul(&self, a: u32, b: u32) -> u32 {
mul_mod(a, b, self._m, self.im)
}
}
/// Calculates `a * b % m`.
///
/// * `a` `0 <= a < m`
/// * `b` `0 <= b < m`
/// * `m` `1 <= m <= 2^31`
/// * `im` = ceil(2^64 / `m`)
#[allow(clippy::many_single_char_names)]
pub(crate) fn mul_mod(a: u32, b: u32, m: u32, im: u64) -> u32 {
// [1] m = 1
// a = b = im = 0, so okay
// [2] m >= 2
// im = ceil(2^64 / m)
// -> im * m = 2^64 + r (0 <= r < m)
// let z = a*b = c*m + d (0 <= c, d < m)
// a*b * im = (c*m + d) * im = c*(im*m) + d*im = c*2^64 + c*r + d*im
// c*r + d*im < m * m + m * im < m * m + 2^64 + m <= 2^64 + m * (m + 1) < 2^64 * 2
// ((ab * im) >> 64) == c or c + 1
let mut z = a as u64;
z *= b as u64;
let x = (((z as u128) * (im as u128)) >> 64) as u64;
let mut v = z.wrapping_sub(x.wrapping_mul(m as u64)) as u32;
if m <= v {
v = v.wrapping_add(m);
}
v
}
/// # Parameters
/// * `n` `0 <= n`
/// * `m` `1 <= m`
///
/// # Returns
/// `(x ** n) % m`
/* const */
#[allow(clippy::many_single_char_names)]
pub(crate) fn pow_mod(x: i64, mut n: i64, m: i32) -> i64 {
if m == 1 {
return 0;
}
let _m = m as u32;
let mut r: u64 = 1;
let mut y: u64 = safe_mod(x, m as i64) as u64;
while n != 0 {
if (n & 1) > 0 {
r = (r * y) % (_m as u64);
}
y = (y * y) % (_m as u64);
n >>= 1;
}
r as i64
}
/// Reference:
/// M. Forisek and J. Jancina,
/// Fast Primality Testing for Integers That Fit into a Machine Word
///
/// # Parameters
/// * `n` `0 <= n`
/* const */
pub(crate) fn is_prime(n: i32) -> bool {
let n = n as i64;
match n {
_ if n <= 1 => return false,
2 | 7 | 61 => return true,
_ if n % 2 == 0 => return false,
_ => {}
}
let mut d = n - 1;
while d % 2 == 0 {
d /= 2;
}
for &a in &[2, 7, 61] {
let mut t = d;
let mut y = pow_mod(a, t, n as i32);
while t != n - 1 && y != 1 && y != n - 1 {
y = y * y % n;
t <<= 1;
}
if y != n - 1 && t % 2 == 0 {
return false;
}
}
true
}
// omitted
// template <int n> constexpr bool is_prime = is_prime_constexpr(n);
/// # Parameters
/// * `b` `1 <= b`
///
/// # Returns
/// (g, x) s.t. g = gcd(a, b), xa = g (mod b), 0 <= x < b/g
/* const */
#[allow(clippy::many_single_char_names)]
pub(crate) fn inv_gcd(a: i64, b: i64) -> (i64, i64) {
let a = safe_mod(a, b);
if a == 0 {
return (b, 0);
}
// Contracts:
// [1] s - m0 * a = 0 (mod b)
// [2] t - m1 * a = 0 (mod b)
// [3] s * |m1| + t * |m0| <= b
let mut s = b;
let mut t = a;
let mut m0 = 0;
let mut m1 = 1;
while t != 0 {
let u = s / t;
s -= t * u;
m0 -= m1 * u; // |m1 * u| <= |m1| * s <= b
// [3]:
// (s - t * u) * |m1| + t * |m0 - m1 * u|
// <= s * |m1| - t * u * |m1| + t * (|m0| + |m1| * u)
// = s * |m1| + t * |m0| <= b
swap(&mut s, &mut t);
swap(&mut m0, &mut m1);
}
// by [3]: |m0| <= b/g
// by g != b: |m0| < b/g
if m0 < 0 {
m0 += b / s;
}
(s, m0)
}
/// Compile time (currently not) primitive root
/// @param m must be prime
/// @return primitive root (and minimum in now)
/* const */
pub(crate) fn primitive_root(m: i32) -> i32 {
match m {
2 => return 1,
167_772_161 => return 3,
469_762_049 => return 3,
754_974_721 => return 11,
998_244_353 => return 3,
_ => {}
}
let mut divs = [0; 20];
divs[0] = 2;
let mut cnt = 1;
let mut x = (m - 1) / 2;
while x % 2 == 0 {
x /= 2;
}
for i in (3..i32::MAX).step_by(2) {
if i as i64 * i as i64 > x as i64 {
break;
}
if x % i == 0 {
divs[cnt] = i;
cnt += 1;
while x % i == 0 {
x /= i;
}
}
}
if x > 1 {
divs[cnt] = x;
cnt += 1;
}
let mut g = 2;
loop {
if (0..cnt).all(|i| pow_mod(g, ((m - 1) / divs[i]) as i64, m) != 1) {
break g as i32;
}
g += 1;
}
}
// omitted
// template <int m> constexpr int primitive_root = primitive_root_constexpr(m);
/// # Arguments
/// * `n` `n < 2^32`
/// * `m` `1 <= m < 2^32`
///
/// # Returns
/// `sum_{i=0}^{n-1} floor((ai + b) / m) (mod 2^64)`
/* const */
#[allow(clippy::many_single_char_names)]
pub(crate) fn floor_sum_unsigned(
mut n: W<u64>,
mut m: W<u64>,
mut a: W<u64>,
mut b: W<u64>,
) -> W<u64> {
let mut ans = W(0);
loop {
if a >= m {
if n > W(0) {
ans += n * (n - W(1)) / W(2) * (a / m);
}
a %= m;
}
if b >= m {
ans += n * (b / m);
b %= m;
}
let y_max = a * n + b;
if y_max < m {
break;
}
// y_max < m * (n + 1)
// floor(y_max / m) <= n
n = y_max / m;
b = y_max % m;
std::mem::swap(&mut m, &mut a);
}
ans
}
}
use std::{
cell::RefCell,
convert::TryInto as _,
hash::{Hash, Hasher},
ops::Neg,
str::FromStr,
sync::atomic::{self, AtomicU32, AtomicU64},
thread::LocalKey,
};
pub type ModInt1000000007 = StaticModInt<Mod1000000007>;
pub type ModInt998244353 = StaticModInt<Mod998244353>;
pub type ModInt = DynamicModInt<DefaultId>;
#[derive(Copy, Clone, Eq, PartialEq)]
#[repr(transparent)]
pub struct StaticModInt<M> {
val: u32,
phantom: PhantomData<fn() -> M>,
}
impl<M: Modulus> StaticModInt<M> {
/// Returns the modulus, which is [`<M as Modulus>::VALUE`].
///
/// Corresponds to `atcoder::static_modint::mod` in the original ACL.
///
/// # Example
///
/// ```
/// use ac_library::ModInt1000000007 as Mint;
///
/// assert_eq!(1_000_000_007, Mint::modulus());
/// ```
///
/// [`<M as Modulus>::VALUE`]: ../trait.Modulus.html#associatedconstant.VALUE
#[inline(always)]
pub fn modulus() -> u32 {
M::VALUE
}
/// Creates a new `StaticModInt`.
///
/// Takes [any primitive integer].
///
/// Corresponds to the constructor of `atcoder::static_modint` in the original ACL.
///
/// [any primitive integer]: ../trait.RemEuclidU32.html
#[inline]
pub fn new<T: RemEuclidU32>(val: T) -> Self {
Self::raw(val.rem_euclid_u32(M::VALUE))
}
/// Constructs a `StaticModInt` from a `val < Self::modulus()` without checking it.
///
/// Corresponds to `atcoder::static_modint::raw` in the original ACL.
///
/// # Constraints
///
/// - `val` is less than `Self::modulus()`
///
/// See [`ModIntBase::raw`] for more more details.
///
/// [`ModIntBase::raw`]: ./trait.ModIntBase.html#tymethod.raw
#[inline]
pub fn raw(val: u32) -> Self {
Self {
val,
phantom: PhantomData,
}
}
/// Returns the representative.
///
/// Corresponds to `atcoder::static_modint::val` in the original ACL.
#[inline]
pub fn val(self) -> u32 {
self.val
}
/// Returns `self` to the power of `n`.
///
/// Corresponds to `atcoder::static_modint::pow` in the original ACL.
#[inline]
pub fn pow(self, n: u64) -> Self {
<Self as ModIntBase>::pow(self, n)
}
/// Returns the multiplicative inverse of `self`.
///
/// Corresponds to `atcoder::static_modint::inv` in the original ACL.
///
/// # Panics
///
/// Panics if the multiplicative inverse does not exist.
#[inline]
pub fn inv(self) -> Self {
if M::HINT_VALUE_IS_PRIME {
if self.val() == 0 {
panic!("attempt to divide by zero");
}
self.pow((M::VALUE - 2).into())
} else {
Self::inv_for_non_prime_modulus(self)
}
}
}
/// These methods are implemented for the struct.
/// You don't need to `use` `ModIntBase` to call methods of `StaticModInt`.
impl<M: Modulus> ModIntBase for StaticModInt<M> {
#[inline(always)]
fn modulus() -> u32 {
Self::modulus()
}
#[inline]
fn raw(val: u32) -> Self {
Self::raw(val)
}
#[inline]
fn val(self) -> u32 {
self.val()
}
#[inline]
fn inv(self) -> Self {
self.inv()
}
}
/// Represents a modulus.
///
/// # Example
///
/// ```
/// macro_rules! modulus {
/// ($($name:ident($value:expr, $is_prime:expr)),*) => {
/// $(
/// #[derive(Copy, Clone, Eq, PartialEq)]
/// enum $name {}
///
/// impl ac_library::modint::Modulus for $name {
/// const VALUE: u32 = $value;
/// const HINT_VALUE_IS_PRIME: bool = $is_prime;
///
/// fn butterfly_cache() -> &'static ::std::thread::LocalKey<::std::cell::RefCell<::std::option::Option<ac_library::modint::ButterflyCache<Self>>>> {
/// thread_local! {
/// static BUTTERFLY_CACHE: ::std::cell::RefCell<::std::option::Option<ac_library::modint::ButterflyCache<$name>>> = ::std::default::Default::default();
/// }
/// &BUTTERFLY_CACHE
/// }
/// }
/// )*
/// };
/// }
///
/// use ac_library::StaticModInt;
///
/// modulus!(Mod101(101, true), Mod103(103, true));
///
/// type Z101 = StaticModInt<Mod101>;
/// type Z103 = StaticModInt<Mod103>;
///
/// assert_eq!(Z101::new(101), Z101::new(0));
/// assert_eq!(Z103::new(103), Z103::new(0));
/// ```
pub trait Modulus: 'static + Copy + Eq {
const VALUE: u32;
const HINT_VALUE_IS_PRIME: bool;
fn butterfly_cache() -> &'static LocalKey<RefCell<Option<ButterflyCache<Self>>>>;
}
/// Represents $1000000007$.
#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
pub enum Mod1000000007 {}
impl Modulus for Mod1000000007 {
const VALUE: u32 = 1_000_000_007;
const HINT_VALUE_IS_PRIME: bool = true;
fn butterfly_cache() -> &'static LocalKey<RefCell<Option<ButterflyCache<Self>>>> {
thread_local! {
static BUTTERFLY_CACHE: RefCell<Option<ButterflyCache<Mod1000000007>>> = RefCell::default();
}
&BUTTERFLY_CACHE
}
}
/// Represents $998244353$.
#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
pub enum Mod998244353 {}
impl Modulus for Mod998244353 {
const VALUE: u32 = 998_244_353;
const HINT_VALUE_IS_PRIME: bool = true;
fn butterfly_cache() -> &'static LocalKey<RefCell<Option<ButterflyCache<Self>>>> {
thread_local! {
static BUTTERFLY_CACHE: RefCell<Option<ButterflyCache<Mod998244353>>> = RefCell::default();
}
&BUTTERFLY_CACHE
}
}
#[allow(unused)]
/// Cache for butterfly operations.
pub struct ButterflyCache<M> {
pub(crate) sum_e: Vec<StaticModInt<M>>,
pub(crate) sum_ie: Vec<StaticModInt<M>>,
}
/// Represents $\mathbb{Z}/m\mathbb{Z}$ where $m$ is a dynamic value.
///
/// Corresponds to `atcoder::dynamic_modint` in the original ACL.
///
/// # Example
///
/// ```
/// use ac_library::ModInt as Mint;
/// use proconio::{input, source::once::OnceSource};
///
/// input! {
/// from OnceSource::from("3 3 7\n"),
/// a: u32,
/// b: u32,
/// m: u32,
/// }
///
/// Mint::set_modulus(m);
/// let a = Mint::new(a);
/// let b = Mint::new(b);
///
/// println!("{}", a * b); // `2`
/// ```
#[derive(Copy, Clone, Eq, PartialEq)]
#[repr(transparent)]
pub struct DynamicModInt<I> {
val: u32,
phantom: PhantomData<fn() -> I>,
}
impl<I: Id> DynamicModInt<I> {
/// Returns the modulus.
///
/// Corresponds to `atcoder::dynamic_modint::mod` in the original ACL.
///
/// # Example
///
/// ```
/// use ac_library::ModInt as Mint;
///
/// assert_eq!(998_244_353, Mint::modulus()); // default modulus
/// ```
#[inline]
pub fn modulus() -> u32 {
I::companion_barrett().umod()
}
/// Sets a modulus.
///
/// Corresponds to `atcoder::dynamic_modint::set_mod` in the original ACL.
///
/// # Constraints
///
/// - This function must be called earlier than any other operation of `Self`.
///
/// # Example
///
/// ```
/// use ac_library::ModInt as Mint;
///
/// Mint::set_modulus(7);
/// assert_eq!(7, Mint::modulus());
/// ```
#[inline]
pub fn set_modulus(modulus: u32) {
if modulus == 0 {
panic!("the modulus must not be 0");
}
I::companion_barrett().update(modulus);
}
/// Creates a new `DynamicModInt`.
///
/// Takes [any primitive integer].
///
/// Corresponds to the constructor of `atcoder::dynamic_modint` in the original ACL.
///
/// [any primitive integer]: ../trait.RemEuclidU32.html
#[inline]
pub fn new<T: RemEuclidU32>(val: T) -> Self {
<Self as ModIntBase>::new(val)
}
/// Constructs a `DynamicModInt` from a `val < Self::modulus()` without checking it.
///
/// Corresponds to `atcoder::dynamic_modint::raw` in the original ACL.
///
/// # Constraints
///
/// - `val` is less than `Self::modulus()`
///
/// See [`ModIntBase::raw`] for more more details.
///
/// [`ModIntBase::raw`]: ./trait.ModIntBase.html#tymethod.raw
#[inline]
pub fn raw(val: u32) -> Self {
Self {
val,
phantom: PhantomData,
}
}
/// Returns the representative.
///
/// Corresponds to `atcoder::static_modint::val` in the original ACL.
#[inline]
pub fn val(self) -> u32 {
self.val
}
/// Returns `self` to the power of `n`.
///
/// Corresponds to `atcoder::dynamic_modint::pow` in the original ACL.
#[inline]
pub fn pow(self, n: u64) -> Self {
<Self as ModIntBase>::pow(self, n)
}
/// Returns the multiplicative inverse of `self`.
///
/// Corresponds to `atcoder::dynamic_modint::inv` in the original ACL.
///
/// # Panics
///
/// Panics if the multiplicative inverse does not exist.
#[inline]
pub fn inv(self) -> Self {
Self::inv_for_non_prime_modulus(self)
}
}
/// These methods are implemented for the struct.
/// You don't need to `use` `ModIntBase` to call methods of `DynamicModInt`.
impl<I: Id> ModIntBase for DynamicModInt<I> {
#[inline]
fn modulus() -> u32 {
Self::modulus()
}
#[inline]
fn raw(val: u32) -> Self {
Self::raw(val)
}
#[inline]
fn val(self) -> u32 {
self.val()
}
#[inline]
fn inv(self) -> Self {
self.inv()
}
}
pub trait Id: 'static + Copy + Eq {
fn companion_barrett() -> &'static Barrett;
}
#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
pub enum DefaultId {}
impl Id for DefaultId {
fn companion_barrett() -> &'static Barrett {
static BARRETT: Barrett = Barrett::default();
&BARRETT
}
}
/// Pair of $m$ and $\lceil 2^{64}/m \rceil$.
pub struct Barrett {
m: AtomicU32,
im: AtomicU64,
}
impl Barrett {
/// Creates a new `Barrett`.
#[inline]
pub const fn new(m: u32) -> Self {
Self {
m: AtomicU32::new(m),
im: AtomicU64::new((-1i64 as u64 / m as u64).wrapping_add(1)),
}
}
#[inline]
const fn default() -> Self {
Self::new(998_244_353)
}
#[inline]
fn update(&self, m: u32) {
let im = (-1i64 as u64 / m as u64).wrapping_add(1);
self.m.store(m, atomic::Ordering::SeqCst);
self.im.store(im, atomic::Ordering::SeqCst);
}
#[inline]
fn umod(&self) -> u32 {
self.m.load(atomic::Ordering::SeqCst)
}
#[inline]
fn mul(&self, a: u32, b: u32) -> u32 {
let m = self.m.load(atomic::Ordering::SeqCst);
let im = self.im.load(atomic::Ordering::SeqCst);
internal_math::mul_mod(a, b, m, im)
}
}
impl Default for Barrett {
#[inline]
fn default() -> Self {
Self::default()
}
}
/// A trait for [`StaticModInt`] and [`DynamicModInt`].
///
/// Corresponds to `atcoder::internal::modint_base` in the original ACL.
///
/// [`StaticModInt`]: ../struct.StaticModInt.html
/// [`DynamicModInt`]: ../struct.DynamicModInt.html
pub trait ModIntBase:
Default
+ FromStr
+ From<i8>
+ From<i16>
+ From<i32>
+ From<i64>
+ From<i128>
+ From<isize>
+ From<u8>
+ From<u16>
+ From<u32>
+ From<u64>
+ From<u128>
+ From<usize>
+ Copy
+ Eq
+ Hash
+ fmt::Display
+ fmt::Debug
+ Neg<Output = Self>
+ Add<Output = Self>
+ Sub<Output = Self>
+ Mul<Output = Self>
+ Div<Output = Self>
+ AddAssign
+ SubAssign
+ MulAssign
+ DivAssign
{
/// Returns the modulus.
///
/// Corresponds to `atcoder::static_modint::mod` and `atcoder::dynamic_modint::mod` in the original ACL.
///
/// # Example
///
/// ```
/// use ac_library::modint::ModIntBase;
///
/// fn f<Z: ModIntBase>() {
/// let _: u32 = Z::modulus();
/// }
/// ```
fn modulus() -> u32;
/// Constructs a `Self` from a `val < Self::modulus()` without checking it.
///
/// Corresponds to `atcoder::static_modint::raw` and `atcoder::dynamic_modint::raw` in the original ACL.
///
/// # Constraints
///
/// - `val` is less than `Self::modulus()`
///
/// **Note that all operations assume that inner values are smaller than the modulus.**
/// If `val` is greater than or equal to `Self::modulus()`, the behaviors are not defined.
///
/// ```should_panic
/// use ac_library::ModInt1000000007 as Mint;
///
/// let x = Mint::raw(1_000_000_007);
/// let y = x + x;
/// assert_eq!(0, y.val());
/// ```
///
/// ```text
/// thread 'main' panicked at 'assertion failed: `(left == right)`
/// left: `0`,
/// right: `1000000007`', src/modint.rs:8:1
/// note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
/// ```
///
/// # Example
///
/// ```
/// use ac_library::modint::ModIntBase;
///
/// fn f<Z: ModIntBase>() -> Z {
/// debug_assert!(Z::modulus() >= 100);
///
/// let mut acc = Z::new(0);
/// for i in 0..100 {
/// if i % 3 == 0 {
/// // I know `i` is smaller than the modulus!
/// acc += Z::raw(i);
/// }
/// }
/// acc
/// }
/// ```
fn raw(val: u32) -> Self;
/// Returns the representative.
///
/// Corresponds to `atcoder::static_modint::val` and `atcoder::dynamic_modint::val` in the original ACL.
///
/// # Example
///
/// ```
/// use ac_library::modint::ModIntBase;
///
/// fn f<Z: ModIntBase>(x: Z) {
/// let _: u32 = x.val();
/// }
/// ```
fn val(self) -> u32;
/// Returns the multiplicative inverse of `self`.
///
/// Corresponds to `atcoder::static_modint::inv` and `atcoder::dynamic_modint::inv` in the original ACL.
///
/// # Panics
///
/// Panics if the multiplicative inverse does not exist.
///
/// # Example
///
/// ```
/// use ac_library::modint::ModIntBase;
///
/// fn f<Z: ModIntBase>(x: Z) {
/// let _: Z = x.inv();
/// }
/// ```
fn inv(self) -> Self;
/// Creates a new `Self`.
///
/// Takes [any primitive integer].
///
/// # Example
///
/// ```
/// use ac_library::modint::ModIntBase;
///
/// fn f<Z: ModIntBase>() {
/// let _ = Z::new(1u32);
/// let _ = Z::new(1usize);
/// let _ = Z::new(-1i64);
/// }
/// ```
///
/// [any primitive integer]: ../trait.RemEuclidU32.html
#[inline]
fn new<T: RemEuclidU32>(val: T) -> Self {
Self::raw(val.rem_euclid_u32(Self::modulus()))
}
/// Returns `self` to the power of `n`.
///
/// Corresponds to `atcoder::static_modint::pow` and `atcoder::dynamic_modint::pow` in the original ACL.
///
/// # Example
///
/// ```
/// use ac_library::modint::ModIntBase;
///
/// fn f<Z: ModIntBase>() {
/// let _: Z = Z::new(2).pow(3);
/// }
/// ```
#[inline]
fn pow(self, mut n: u64) -> Self {
let mut x = self;
let mut r = Self::raw(u32::from(Self::modulus() > 1));
while n > 0 {
if n & 1 == 1 {
r *= x;
}
x *= x;
n >>= 1;
}
r
}
}
/// A trait for `{StaticModInt, DynamicModInt, ModIntBase}::new`.
pub trait RemEuclidU32 {
/// Calculates `self` $\bmod$ `modulus` losslessly.
fn rem_euclid_u32(self, modulus: u32) -> u32;
}
macro_rules! impl_rem_euclid_u32_for_small_signed {
($($ty:tt),*) => {
$(
impl RemEuclidU32 for $ty {
#[inline]
fn rem_euclid_u32(self, modulus: u32) -> u32 {
(self as i64).rem_euclid(i64::from(modulus)) as _
}
}
)*
}
}
impl_rem_euclid_u32_for_small_signed!(i8, i16, i32, i64, isize);
impl RemEuclidU32 for i128 {
#[inline]
fn rem_euclid_u32(self, modulus: u32) -> u32 {
self.rem_euclid(i128::from(modulus)) as _
}
}
macro_rules! impl_rem_euclid_u32_for_small_unsigned {
($($ty:tt),*) => {
$(
impl RemEuclidU32 for $ty {
#[inline]
fn rem_euclid_u32(self, modulus: u32) -> u32 {
self as u32 % modulus
}
}
)*
}
}
macro_rules! impl_rem_euclid_u32_for_large_unsigned {
($($ty:tt),*) => {
$(
impl RemEuclidU32 for $ty {
#[inline]
fn rem_euclid_u32(self, modulus: u32) -> u32 {
(self % (modulus as $ty)) as _
}
}
)*
}
}
impl_rem_euclid_u32_for_small_unsigned!(u8, u16, u32);
impl_rem_euclid_u32_for_large_unsigned!(u64, u128);
#[cfg(target_pointer_width = "32")]
impl_rem_euclid_u32_for_small_unsigned!(usize);
#[cfg(target_pointer_width = "64")]
impl_rem_euclid_u32_for_large_unsigned!(usize);
trait InternalImplementations: ModIntBase {
#[inline]
fn inv_for_non_prime_modulus(this: Self) -> Self {
let (gcd, x) = internal_math::inv_gcd(this.val().into(), Self::modulus().into());
if gcd != 1 {
panic!("the multiplicative inverse does not exist");
}
Self::new(x)
}
#[inline]
fn default_impl() -> Self {
Self::raw(0)
}
#[inline]
fn from_str_impl(s: &str) -> Result<Self, Infallible> {
Ok(s.parse::<i64>()
.map(Self::new)
.unwrap_or_else(|_| todo!("parsing as an arbitrary precision integer?")))
}
#[inline]
fn hash_impl(this: &Self, state: &mut impl Hasher) {
this.val().hash(state)
}
#[inline]
fn display_impl(this: &Self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Display::fmt(&this.val(), f)
}
#[inline]
fn debug_impl(this: &Self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&this.val(), f)
}
#[inline]
fn neg_impl(this: Self) -> Self {
Self::sub_impl(Self::raw(0), this)
}
#[inline]
fn add_impl(lhs: Self, rhs: Self) -> Self {
let modulus = Self::modulus();
let mut val = lhs.val() + rhs.val();
if val >= modulus {
val -= modulus;
}
Self::raw(val)
}
#[inline]
fn sub_impl(lhs: Self, rhs: Self) -> Self {
let modulus = Self::modulus();
let mut val = lhs.val().wrapping_sub(rhs.val());
if val >= modulus {
val = val.wrapping_add(modulus)
}
Self::raw(val)
}
fn mul_impl(lhs: Self, rhs: Self) -> Self;
#[inline]
fn div_impl(lhs: Self, rhs: Self) -> Self {
Self::mul_impl(lhs, rhs.inv())
}
}
impl<M: Modulus> InternalImplementations for StaticModInt<M> {
#[inline]
fn mul_impl(lhs: Self, rhs: Self) -> Self {
Self::raw((u64::from(lhs.val()) * u64::from(rhs.val()) % u64::from(M::VALUE)) as u32)
}
}
impl<I: Id> InternalImplementations for DynamicModInt<I> {
#[inline]
fn mul_impl(lhs: Self, rhs: Self) -> Self {
Self::raw(I::companion_barrett().mul(lhs.val, rhs.val))
}
}
macro_rules! impl_basic_traits {
() => {};
(impl <$generic_param:ident : $generic_param_bound:tt> _ for $self:ty; $($rest:tt)*) => {
impl <$generic_param: $generic_param_bound> Default for $self {
#[inline]
fn default() -> Self {
Self::default_impl()
}
}
impl <$generic_param: $generic_param_bound> FromStr for $self {
type Err = Infallible;
#[inline]
fn from_str(s: &str) -> Result<Self, Infallible> {
Self::from_str_impl(s)
}
}
impl<$generic_param: $generic_param_bound, V: RemEuclidU32> From<V> for $self {
#[inline]
fn from(from: V) -> Self {
Self::new(from)
}
}
#[allow(clippy::derived_hash_with_manual_eq)]
impl<$generic_param: $generic_param_bound> Hash for $self {
#[inline]
fn hash<H: Hasher>(&self, state: &mut H) {
Self::hash_impl(self, state)
}
}
impl<$generic_param: $generic_param_bound> fmt::Display for $self {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Self::display_impl(self, f)
}
}
impl<$generic_param: $generic_param_bound> fmt::Debug for $self {
#[inline]
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
Self::debug_impl(self, f)
}
}
impl<$generic_param: $generic_param_bound> Neg for $self {
type Output = $self;
#[inline]
fn neg(self) -> $self {
Self::neg_impl(self)
}
}
impl<$generic_param: $generic_param_bound> Neg for &'_ $self {
type Output = $self;
#[inline]
fn neg(self) -> $self {
<$self>::neg_impl(*self)
}
}
impl_basic_traits!($($rest)*);
};
}
impl_basic_traits! {
impl <M: Modulus> _ for StaticModInt<M> ;
impl <I: Id > _ for DynamicModInt<I>;
}
macro_rules! impl_bin_ops {
() => {};
(for<$($generic_param:ident : $generic_param_bound:tt),*> <$lhs_ty:ty> ~ <$rhs_ty:ty> -> $output:ty { { $lhs_body:expr } ~ { $rhs_body:expr } } $($rest:tt)*) => {
impl <$($generic_param: $generic_param_bound),*> Add<$rhs_ty> for $lhs_ty {
type Output = $output;
#[inline]
fn add(self, rhs: $rhs_ty) -> $output {
<$output>::add_impl(apply($lhs_body, self), apply($rhs_body, rhs))
}
}
impl <$($generic_param: $generic_param_bound),*> Sub<$rhs_ty> for $lhs_ty {
type Output = $output;
#[inline]
fn sub(self, rhs: $rhs_ty) -> $output {
<$output>::sub_impl(apply($lhs_body, self), apply($rhs_body, rhs))
}
}
impl <$($generic_param: $generic_param_bound),*> Mul<$rhs_ty> for $lhs_ty {
type Output = $output;
#[inline]
fn mul(self, rhs: $rhs_ty) -> $output {
<$output>::mul_impl(apply($lhs_body, self), apply($rhs_body, rhs))
}
}
impl <$($generic_param: $generic_param_bound),*> Div<$rhs_ty> for $lhs_ty {
type Output = $output;
#[inline]
fn div(self, rhs: $rhs_ty) -> $output {
<$output>::div_impl(apply($lhs_body, self), apply($rhs_body, rhs))
}
}
impl_bin_ops!($($rest)*);
};
}
macro_rules! impl_assign_ops {
() => {};
(for<$($generic_param:ident : $generic_param_bound:tt),*> <$lhs_ty:ty> ~= <$rhs_ty:ty> { _ ~= { $rhs_body:expr } } $($rest:tt)*) => {
impl <$($generic_param: $generic_param_bound),*> AddAssign<$rhs_ty> for $lhs_ty {
#[inline]
fn add_assign(&mut self, rhs: $rhs_ty) {
*self = *self + apply($rhs_body, rhs);
}
}
impl <$($generic_param: $generic_param_bound),*> SubAssign<$rhs_ty> for $lhs_ty {
#[inline]
fn sub_assign(&mut self, rhs: $rhs_ty) {
*self = *self - apply($rhs_body, rhs);
}
}
impl <$($generic_param: $generic_param_bound),*> MulAssign<$rhs_ty> for $lhs_ty {
#[inline]
fn mul_assign(&mut self, rhs: $rhs_ty) {
*self = *self * apply($rhs_body, rhs);
}
}
impl <$($generic_param: $generic_param_bound),*> DivAssign<$rhs_ty> for $lhs_ty {
#[inline]
fn div_assign(&mut self, rhs: $rhs_ty) {
*self = *self / apply($rhs_body, rhs);
}
}
impl_assign_ops!($($rest)*);
};
}
#[inline]
fn apply<F: FnOnce(X) -> O, X, O>(f: F, x: X) -> O {
f(x)
}
impl_bin_ops! {
for<M: Modulus> <StaticModInt<M> > ~ <StaticModInt<M> > -> StaticModInt<M> { { |x| x } ~ { |x| x } }
for<M: Modulus> <StaticModInt<M> > ~ <&'_ StaticModInt<M> > -> StaticModInt<M> { { |x| x } ~ { |&x| x } }
for<M: Modulus> <&'_ StaticModInt<M> > ~ <StaticModInt<M> > -> StaticModInt<M> { { |&x| x } ~ { |x| x } }
for<M: Modulus> <&'_ StaticModInt<M> > ~ <&'_ StaticModInt<M> > -> StaticModInt<M> { { |&x| x } ~ { |&x| x } }
for<I: Id > <DynamicModInt<I> > ~ <DynamicModInt<I> > -> DynamicModInt<I> { { |x| x } ~ { |x| x } }
for<I: Id > <DynamicModInt<I> > ~ <&'_ DynamicModInt<I>> -> DynamicModInt<I> { { |x| x } ~ { |&x| x } }
for<I: Id > <&'_ DynamicModInt<I>> ~ <DynamicModInt<I> > -> DynamicModInt<I> { { |&x| x } ~ { |x| x } }
for<I: Id > <&'_ DynamicModInt<I>> ~ <&'_ DynamicModInt<I>> -> DynamicModInt<I> { { |&x| x } ~ { |&x| x } }
for<M: Modulus, T: RemEuclidU32> <StaticModInt<M> > ~ <T> -> StaticModInt<M> { { |x| x } ~ { StaticModInt::<M>::new } }
for<I: Id , T: RemEuclidU32> <DynamicModInt<I> > ~ <T> -> DynamicModInt<I> { { |x| x } ~ { DynamicModInt::<I>::new } }
}
impl_assign_ops! {
for<M: Modulus> <StaticModInt<M> > ~= <StaticModInt<M> > { _ ~= { |x| x } }
for<M: Modulus> <StaticModInt<M> > ~= <&'_ StaticModInt<M> > { _ ~= { |&x| x } }
for<I: Id > <DynamicModInt<I>> ~= <DynamicModInt<I> > { _ ~= { |x| x } }
for<I: Id > <DynamicModInt<I>> ~= <&'_ DynamicModInt<I>> { _ ~= { |&x| x } }
for<M: Modulus, T: RemEuclidU32> <StaticModInt<M> > ~= <T> { _ ~= { StaticModInt::<M>::new } }
for<I: Id, T: RemEuclidU32> <DynamicModInt<I>> ~= <T> { _ ~= { DynamicModInt::<I>::new } }
}
macro_rules! impl_folding {
() => {};
(impl<$generic_param:ident : $generic_param_bound:tt> $trait:ident<_> for $self:ty { fn $method:ident(_) -> _ { _($unit:expr, $op:expr) } } $($rest:tt)*) => {
impl<$generic_param: $generic_param_bound> $trait<Self> for $self {
#[inline]
fn $method<S>(iter: S) -> Self
where
S: Iterator<Item = Self>,
{
iter.fold($unit, $op)
}
}
impl<'a, $generic_param: $generic_param_bound> $trait<&'a Self> for $self {
#[inline]
fn $method<S>(iter: S) -> Self
where
S: Iterator<Item = &'a Self>,
{
iter.fold($unit, $op)
}
}
impl_folding!($($rest)*);
};
}
impl_folding! {
impl<M: Modulus> Sum<_> for StaticModInt<M> { fn sum(_) -> _ { _(Self::raw(0), Add::add) } }
impl<M: Modulus> Product<_> for StaticModInt<M> { fn product(_) -> _ { _(Self::raw(u32::from(Self::modulus() > 1)), Mul::mul) } }
impl<I: Id > Sum<_> for DynamicModInt<I> { fn sum(_) -> _ { _(Self::raw(0), Add::add) } }
impl<I: Id > Product<_> for DynamicModInt<I> { fn product(_) -> _ { _(Self::raw(u32::from(Self::modulus() > 1)), Mul::mul) } }
}
/// A Disjoint set union (DSU) with union by size and path compression.
///
/// See: [Zvi Galil and Giuseppe F. Italiano, Data structures and algorithms for disjoint set union problems](https://core.ac.uk/download/pdf/161439519.pdf)
///
/// In the following documentation, let $n$ be the size of the DSU.
///
/// # Example
///
/// ```
/// use ac_library::Dsu;
/// use proconio::{input, source::once::OnceSource};
///
/// input! {
/// from OnceSource::from(
/// "5\n\
/// 3\n\
/// 0 1\n\
/// 2 3\n\
/// 3 4\n",
/// ),
/// n: usize,
/// abs: [(usize, usize)],
/// }
///
/// let mut dsu = Dsu::new(n);
/// for (a, b) in abs {
/// dsu.merge(a, b);
/// }
///
/// assert!(dsu.same(0, 1));
/// assert!(!dsu.same(1, 2));
/// assert!(dsu.same(2, 4));
///
/// assert_eq!(
/// dsu.groups()
/// .into_iter()
/// .map(|mut group| {
/// group.sort_unstable();
/// group
/// })
/// .collect::<Vec<_>>(),
/// [&[0, 1][..], &[2, 3, 4][..]],
/// );
/// ```
#[derive(Clone, Debug)]
pub struct Dsu {
n: usize,
// root node: -1 * component size
// otherwise: parent
parent_or_size: Vec<i32>,
}
impl Dsu {
/// Creates a new `Dsu`.
///
/// # Constraints
///
/// - $0 \leq n \leq 10^8$
///
/// # Complexity
///
/// - $O(n)$
pub fn new(size: usize) -> Self {
Self {
n: size,
parent_or_size: vec![-1; size],
}
}
// `\textsc` does not work in KaTeX
/// Performs the Uɴɪᴏɴ operation.
///
/// # Constraints
///
/// - $0 \leq a < n$
/// - $0 \leq b < n$
///
/// # Panics
///
/// Panics if the above constraints are not satisfied.
///
/// # Complexity
///
/// - $O(\alpha(n))$ amortized
pub fn merge(&mut self, a: usize, b: usize) -> usize {
assert!(a < self.n);
assert!(b < self.n);
let (mut x, mut y) = (self.leader(a), self.leader(b));
if x == y {
return x;
}
if -self.parent_or_size[x] < -self.parent_or_size[y] {
std::mem::swap(&mut x, &mut y);
}
self.parent_or_size[x] += self.parent_or_size[y];
self.parent_or_size[y] = x as i32;
x
}
/// Returns whether the vertices $a$ and $b$ are in the same connected component.
///
/// # Constraints
///
/// - $0 \leq a < n$
/// - $0 \leq b < n$
///
/// # Panics
///
/// Panics if the above constraint is not satisfied.
///
/// # Complexity
///
/// - $O(\alpha(n))$ amortized
pub fn same(&mut self, a: usize, b: usize) -> bool {
assert!(a < self.n);
assert!(b < self.n);
self.leader(a) == self.leader(b)
}
/// Performs the Fɪɴᴅ operation.
///
/// # Constraints
///
/// - $0 \leq a < n$
///
/// # Panics
///
/// Panics if the above constraint is not satisfied.
///
/// # Complexity
///
/// - $O(\alpha(n))$ amortized
pub fn leader(&mut self, a: usize) -> usize {
assert!(a < self.n);
self._leader(a)
}
/// Returns the size of the connected component that contains the vertex $a$.
///
/// # Constraints
///
/// - $0 \leq a < n$
///
/// # Panics
///
/// Panics if the above constraint is not satisfied.
///
/// # Complexity
///
/// - $O(\alpha(n))$ amortized
pub fn size(&mut self, a: usize) -> usize {
assert!(a < self.n);
let x = self.leader(a);
-self.parent_or_size[x] as usize
}
/// Divides the graph into connected components.
///
/// The result may not be ordered.
///
/// # Complexity
///
/// - $O(n)$
pub fn groups(&mut self) -> Vec<Vec<usize>> {
let mut leader_buf = vec![0; self.n];
let mut group_size = vec![0; self.n];
for i in 0..self.n {
leader_buf[i] = self.leader(i);
group_size[leader_buf[i]] += 1;
}
let mut result = vec![Vec::new(); self.n];
for i in 0..self.n {
result[i].reserve(group_size[i]);
}
for i in 0..self.n {
result[leader_buf[i]].push(i);
}
result
.into_iter()
.filter(|x| !x.is_empty())
.collect::<Vec<Vec<usize>>>()
}
fn _leader(&mut self, a: usize) -> usize {
if self.parent_or_size[a] < 0 {
return a;
}
self.parent_or_size[a] = self._leader(self.parent_or_size[a] as usize) as i32;
self.parent_or_size[a] as usize
}
}