competitive/num/
integer.rs

1use super::{Bounded, IterScan, One, Zero};
2use std::{
3    convert::TryFrom,
4    fmt::{self, Display},
5    iter::{Product, Sum},
6    ops::{
7        Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div,
8        DivAssign, Mul, MulAssign, Neg, Not, Rem, RemAssign, Shl, ShlAssign, Shr, ShrAssign, Sub,
9        SubAssign,
10    },
11    str::FromStr,
12};
13
14// primitive integer = arithmetic operations + binary represented operation
15// arithmetic operations = integer basic operations + (unsigned operations | signed operations)
16
17/// Trait for basic primitive integer operations.
18pub trait IntBase:
19    Copy
20    + Bounded
21    + Zero
22    + One
23    + Eq
24    + Ord
25    + Default
26    + FromStr
27    + Display
28    + Add<Output = Self>
29    + Sub<Output = Self>
30    + Mul<Output = Self>
31    + Div<Output = Self>
32    + Rem<Output = Self>
33    + AddAssign
34    + SubAssign
35    + MulAssign
36    + DivAssign
37    + RemAssign
38    + Sum
39    + Product
40{
41    type Error;
42    fn div_euclid(self, rhs: Self) -> Self;
43    fn rem_euclid(self, rhs: Self) -> Self;
44    fn pow(self, exp: u32) -> Self;
45    fn from_str_radix(src: &str, radix: u32) -> Result<Self, Self::Error>;
46    fn ilog(self, base: Self) -> u32;
47    fn ilog2(self) -> u32;
48    fn ilog10(self) -> u32;
49}
50macro_rules! impl_int_base {
51    ($($t:ty)*) => {
52        $(
53            impl IntBase for $t {
54                type Error = std::num::ParseIntError;
55                fn div_euclid(self, rhs: Self) -> Self { self.div_euclid(rhs) }
56                fn rem_euclid(self, rhs: Self) -> Self { self.rem_euclid(rhs) }
57                fn pow(self, exp: u32) -> Self { self.pow(exp) }
58                fn from_str_radix(src: &str, radix: u32) -> Result<Self, Self::Error> { Self::from_str_radix(src, radix) }
59                fn ilog(self, base: Self) -> u32 { self.ilog(base) }
60                fn ilog2(self) -> u32 { self.ilog2() }
61                fn ilog10(self) -> u32 { self.ilog10() }
62            }
63        )*
64    };
65}
66impl_int_base!(u8 i8 u16 i16 u32 i32 u64 i64 u128 i128 usize isize);
67
68/// extended_gcd(a,b): ax + by = g = gcd(a,b)
69pub struct ExtendedGcd<T: Signed> {
70    /// gcd
71    pub g: T::Unsigned,
72    pub x: T,
73    pub y: T,
74}
75
76/// Trait for unsigned integer operations.
77pub trait Unsigned: IntBase {
78    type Signed: Signed<Unsigned = Self>;
79    fn signed(self) -> Self::Signed;
80    fn abs_diff(self, other: Self) -> Self;
81    fn next_power_of_two(self) -> Self;
82    fn gcd(self, other: Self) -> Self;
83    fn lcm(self, other: Self) -> Self {
84        if self.is_zero() && other.is_zero() {
85            Self::zero()
86        } else {
87            self / self.gcd(other) * other
88        }
89    }
90    fn mod_inv(self, modulo: Self) -> Self {
91        assert!(
92            !self.is_zero(),
93            "attempt to inverse zero with modulo {}",
94            modulo
95        );
96        let extgcd = self.signed().extgcd(modulo.signed());
97        assert!(
98            extgcd.g.is_one(),
99            "there is no inverse {} modulo {}",
100            self,
101            modulo
102        );
103        extgcd.x.rem_euclid(modulo.signed()).unsigned()
104    }
105    fn mod_add(self, rhs: Self, modulo: Self) -> Self;
106    fn mod_sub(self, rhs: Self, modulo: Self) -> Self;
107    fn mod_mul(self, rhs: Self, modulo: Self) -> Self;
108    fn mod_neg(self, modulo: Self) -> Self {
109        debug_assert!(!modulo.is_zero(), "modulo must be non-zero");
110        debug_assert!(self < modulo, "self must be less than modulo");
111        if self.is_zero() {
112            Self::zero()
113        } else {
114            modulo - self
115        }
116    }
117}
118
119/// Trait for signed integer operations.
120pub trait Signed: IntBase + Neg<Output = Self> {
121    type Unsigned: Unsigned<Signed = Self>;
122    fn unsigned(self) -> Self::Unsigned;
123    fn abs(self) -> Self;
124    fn abs_diff(self, other: Self) -> Self::Unsigned;
125    fn is_negative(self) -> bool;
126    fn is_positive(self) -> bool;
127    fn signum(self) -> Self;
128    fn extgcd(self, other: Self) -> ExtendedGcd<Self> {
129        let (mut a, mut b) = (self, other);
130        let (mut u, mut v, mut x, mut y) = (Self::one(), Self::zero(), Self::zero(), Self::one());
131        while !a.is_zero() {
132            let k = b / a;
133            x -= k * u;
134            y -= k * v;
135            b -= k * a;
136            std::mem::swap(&mut x, &mut u);
137            std::mem::swap(&mut y, &mut v);
138            std::mem::swap(&mut b, &mut a);
139        }
140        if b.is_negative() {
141            b = -b;
142            x = -x;
143            y = -y;
144        }
145        ExtendedGcd {
146            g: b.unsigned(),
147            x,
148            y,
149        }
150    }
151}
152
153macro_rules! impl_unsigned_signed {
154    ($([$($tt:tt)*])*) => {
155        $(impl_unsigned_signed!($($tt)*);)*
156    };
157    ($unsigned:ident $signed:ident $upperty:ident) => {
158        impl_unsigned_signed!(
159            @inner $unsigned $signed
160            fn mod_mul(self, rhs: Self, modulo: Self) -> Self {
161                debug_assert!(!modulo.is_zero(), "modulo must be non-zero");
162                (self as $upperty * rhs as $upperty % modulo as $upperty) as $unsigned
163            }
164        );
165    };
166    (u128 i128) => {
167        impl_unsigned_signed!(
168            @inner u128 i128
169            fn mod_mul(self, rhs: Self, modulo: Self) -> Self {
170                debug_assert!(!modulo.is_zero(), "modulo must be non-zero");
171                const MASK64: u128 = 0xffff_ffff_ffff_ffff;
172                let (au, ad) = (self >> 64, self & MASK64);
173                let (bu, bd) = (rhs >> 64, rhs & MASK64);
174                let p0 = ad * bd % modulo;
175                let p2 = au * bu % modulo;
176                let mut x = [
177                    p0 as u64,
178                    (p0 >> 64) as u64,
179                    p2 as u64,
180                    (p2 >> 64) as u64,
181                ];
182                let p1 = (au * bd % modulo).mod_add(ad * bu % modulo, modulo);
183                let (p1_lo, p1_hi) = ((p1 & MASK64) as u64, (p1 >> 64) as u64);
184                let (s1, c1) = x[1].overflowing_add(p1_lo);
185                x[1] = s1 as u64;
186                let (s2, c2) = x[2].overflowing_add(p1_hi + c1 as u64);
187                x[2] = s2 as u64;
188                let (s3, _) = x[3].overflowing_add(c2 as u64);
189                x[3] = s3 as u64;
190                rem_u256_by_u128(x, modulo)
191            }
192        );
193    };
194    (@inner $unsigned:ident $signed:ident $mod_mul:item) => {
195        impl Unsigned for $unsigned {
196            type Signed = $signed;
197            fn signed(self) -> Self::Signed { self as Self::Signed }
198            fn abs_diff(self, other: Self) -> Self { self.abs_diff(other) }
199            fn next_power_of_two(self) -> Self { self.next_power_of_two() }
200            fn gcd(self, other: Self) -> Self {
201                let (mut a, mut b) = (self, other);
202                if a.is_zero() || b.is_zero() {
203                    return a | b;
204                }
205                let u = a.trailing_zeros();
206                let v = b.trailing_zeros();
207                a >>= u;
208                b >>= v;
209                let k = u.min(v);
210                while a != b {
211                    if a < b {
212                        std::mem::swap(&mut a, &mut b);
213                    }
214                    a -= b;
215                    a >>= a.trailing_zeros();
216                }
217                a << k
218            }
219            fn mod_add(self, rhs: Self, modulo: Self) -> Self {
220                debug_assert!(!modulo.is_zero(), "modulo must be non-zero");
221                debug_assert!(self < modulo, "self must be less than modulo");
222                debug_assert!(rhs < modulo, "rhs must be less than modulo");
223                let s = self.wrapping_add(rhs);
224                if (s < self) || (s >= modulo) { s.wrapping_sub(modulo) } else { s }
225            }
226            fn mod_sub(self, rhs: Self, modulo: Self) -> Self {
227                debug_assert!(!modulo.is_zero(), "modulo must be non-zero");
228                debug_assert!(self < modulo, "self must be less than modulo");
229                debug_assert!(rhs < modulo, "rhs must be less than modulo");
230                let d = self.wrapping_sub(rhs);
231                if self < rhs { d.wrapping_add(modulo) } else { d }
232            }
233            $mod_mul
234        }
235        impl Signed for $signed {
236            type Unsigned = $unsigned;
237            fn unsigned(self) -> Self::Unsigned { self as Self::Unsigned }
238            fn abs_diff(self, other: Self) -> Self::Unsigned { self.abs_diff(other) }
239            fn abs(self) -> Self { self.abs() }
240            fn is_negative(self) -> bool { self.is_negative() }
241            fn is_positive(self) -> bool { self.is_positive() }
242            fn signum(self) -> Self { self.signum() }
243        }
244    };
245}
246impl_unsigned_signed!([u8 i8 u16] [u16 i16 u32] [u32 i32 u64] [u64 i64 u128] [u128 i128] [usize isize u128]);
247
248fn rem_u256_by_u128(u: [u64; 4], v: u128) -> u128 {
249    // FIXME: use carrying_add and carrying_sub when stabilized
250    #[inline(always)]
251    fn sub_with_borrow_u64(lhs: u64, rhs: u64, borrow: bool) -> (u64, bool) {
252        let (res, overflow) = lhs.overflowing_sub(rhs);
253        if borrow {
254            let (res, overflow_borrow) = res.overflowing_sub(1);
255            (res, overflow | overflow_borrow)
256        } else {
257            (res, overflow)
258        }
259    }
260
261    #[inline(always)]
262    fn add_with_carry_u64(lhs: u64, rhs: u64, carry: bool) -> (u64, bool) {
263        let (res, overflow) = lhs.overflowing_add(rhs);
264        if carry {
265            let (res, overflow_carry) = res.overflowing_add(1);
266            (res, overflow | overflow_carry)
267        } else {
268            (res, overflow)
269        }
270    }
271
272    debug_assert!(v != 0);
273    let v_hi = (v >> 64) as u64;
274    if v_hi == 0 {
275        let d = v as u64 as u128;
276        let mut rem: u128 = 0;
277        for &w in u.iter().rev() {
278            rem = (rem << 64 | w as u128) % d;
279        }
280        return rem;
281    }
282
283    let v_lo = v as u64;
284    let v_shift = v_hi.leading_zeros();
285    let (vn1, vn0) = if v_shift == 0 {
286        (v_hi, v_lo)
287    } else {
288        let hi = v_hi << v_shift | v_lo >> (64 - v_shift);
289        let lo = v_lo << v_shift;
290        (hi, lo)
291    };
292
293    let mut un = [0u64; 5];
294    if v_shift == 0 {
295        un[0] = u[0];
296        un[1] = u[1];
297        un[2] = u[2];
298        un[3] = u[3];
299    } else {
300        un[0] = u[0] << v_shift;
301        un[1] = u[1] << v_shift | u[0] >> (64 - v_shift);
302        un[2] = u[2] << v_shift | u[1] >> (64 - v_shift);
303        un[3] = u[3] << v_shift | u[2] >> (64 - v_shift);
304        un[4] = u[3] >> (64 - v_shift);
305    }
306
307    for j in (0..=2).rev() {
308        let num = (un[j + 2] as u128) << 64 | un[j + 1] as u128;
309        let (mut qhat, mut rhat) = if un[j + 2] == vn1 {
310            (u64::MAX, un[j + 1])
311        } else {
312            let d = vn1 as u128;
313            ((num / d) as u64, (num % d) as u64)
314        };
315        while qhat as u128 * vn0 as u128 > (rhat as u128) << 64 | un[j] as u128 {
316            qhat -= 1;
317            let t = rhat as u128 + vn1 as u128;
318            if t >= 1u128 << 64 {
319                break;
320            }
321            rhat = t as u64;
322        }
323
324        let p0 = qhat as u128 * vn0 as u128;
325        let p1 = qhat as u128 * vn1 as u128;
326        let (p0_hi, p0_lo) = ((p0 >> 64) as u64, p0 as u64);
327        let (p1_hi, p1_lo) = ((p1 >> 64) as u64, p1 as u64);
328
329        let (r0, borrow) = sub_with_borrow_u64(un[j], p0_lo, false);
330        un[j] = r0;
331
332        let (r1, borrow1) = sub_with_borrow_u64(un[j + 1], p0_hi, borrow);
333        let (r1, borrow2) = sub_with_borrow_u64(r1, p1_lo, false);
334        let borrow = borrow1 || borrow2;
335        un[j + 1] = r1;
336
337        let (r2, borrow) = sub_with_borrow_u64(un[j + 2], p1_hi, borrow);
338        un[j + 2] = r2;
339
340        if borrow {
341            let (s0, carry) = add_with_carry_u64(un[j], vn0, false);
342            un[j] = s0;
343            let (s1, carry) = add_with_carry_u64(un[j + 1], vn1, carry);
344            un[j + 1] = s1;
345            let (s2, _) = un[j + 2].overflowing_add(carry as u64);
346            un[j + 2] = s2;
347        }
348    }
349
350    ((un[1] as u128) << 64 | un[0] as u128) >> v_hi.leading_zeros()
351}
352
353/// Trait for operations of integer in binary representation.
354pub trait BinaryRepr<Size = u32>:
355    Sized
356    + Not<Output = Self>
357    + BitAnd<Output = Self>
358    + BitOr<Output = Self>
359    + BitXor<Output = Self>
360    + Shl<Size, Output = Self>
361    + Shr<Size, Output = Self>
362    + BitAndAssign
363    + BitOrAssign
364    + BitXorAssign
365    + ShlAssign<Size>
366    + ShrAssign<Size>
367{
368    fn count_ones(self) -> Size;
369    fn count_zeros(self) -> Size;
370    fn leading_ones(self) -> Size;
371    fn leading_zeros(self) -> Size;
372    fn reverse_bits(self) -> Self;
373    fn rotate_left(self, n: Size) -> Self;
374    fn rotate_right(self, n: Size) -> Self;
375    fn swap_bytes(self) -> Self;
376    fn trailing_ones(self) -> Size;
377    fn trailing_zeros(self) -> Size;
378}
379
380macro_rules! impl_binary_repr {
381    ($($t:ty)*) => {
382        $(
383            impl BinaryRepr for $t {
384                fn count_ones(self) -> u32 { self.count_ones() }
385                fn count_zeros(self) -> u32 { self.count_zeros() }
386                fn leading_ones(self) -> u32 { self.leading_ones() }
387                fn leading_zeros(self) -> u32 { self.leading_zeros() }
388                fn reverse_bits(self) -> Self { self.reverse_bits() }
389                fn rotate_left(self, n: u32) -> Self { self.rotate_left(n) }
390                fn rotate_right(self, n: u32) -> Self { self.rotate_right(n) }
391                fn swap_bytes(self) -> Self { self.swap_bytes() }
392                fn trailing_ones(self) -> u32 { self.trailing_ones() }
393                fn trailing_zeros(self) -> u32 { self.trailing_zeros() }
394            }
395        )*
396    };
397}
398impl_binary_repr!(u8 i8 u16 i16 u32 i32 u64 i64 u128 i128 usize isize);
399
400macro_rules! impl_binop {
401    (impl<$T:ident> $Trait:ident $impl:ident for $t:ty) => {
402        impl<$T> $Trait for $t
403        where
404            $T: $Trait<Output = $T>,
405        {
406            type Output = Self;
407            fn $impl(self, rhs: Self) -> Self::Output {
408                Self($Trait::$impl(self.0, rhs.0))
409            }
410        }
411        impl<$T> $Trait<$T> for $t
412        where
413            $T: $Trait<Output = $T>,
414        {
415            type Output = Self;
416            fn $impl(self, rhs: $T) -> Self::Output {
417                Self($Trait::$impl(self.0, rhs))
418            }
419        }
420    };
421}
422macro_rules! impl_opassign {
423    (impl<$T:ident> $Trait:ident $impl:ident for $t:ty) => {
424        impl<$T> $Trait for $t
425        where
426            $T: $Trait,
427        {
428            fn $impl(&mut self, rhs: Self) {
429                $Trait::$impl(&mut self.0, rhs.0)
430            }
431        }
432        impl<$T> $Trait<$T> for $t
433        where
434            $T: $Trait,
435        {
436            fn $impl(&mut self, rhs: $T) {
437                $Trait::$impl(&mut self.0, rhs)
438            }
439        }
440    };
441    (impl<$T:ident> $Trait:ident $impl:ident for $t:ty => $F:ident $f:ident) => {
442        impl<$T> $Trait for $t
443        where
444            $t: $F<Output = $t> + Copy,
445        {
446            fn $impl(&mut self, rhs: Self) {
447                *self = $F::$f(*self, rhs);
448            }
449        }
450        impl<$T> $Trait<$T> for $t
451        where
452            $t: $F<$T, Output = $t> + Copy,
453        {
454            fn $impl(&mut self, rhs: $T) {
455                *self = $F::$f(*self, rhs);
456            }
457        }
458    };
459}
460
461#[derive(Default, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
462#[repr(transparent)]
463/// Wrapper type of arithmetic `saturating_*` operations.
464pub struct Saturating<T>(pub T);
465pub trait Saturatingable: Sized
466where
467    Saturating<Self>: Copy
468        + Bounded
469        + Zero
470        + One
471        + Eq
472        + Ord
473        + Default
474        + FromStr
475        + Display
476        + Add<Output = Saturating<Self>>
477        + Sub<Output = Saturating<Self>>
478        + Mul<Output = Saturating<Self>>
479        + Div<Output = Saturating<Self>>
480        + Rem<Output = Saturating<Self>>
481        + BitAnd<Output = Saturating<Self>>
482        + BitOr<Output = Saturating<Self>>
483        + BitXor<Output = Saturating<Self>>
484        + Shl<u32, Output = Saturating<Self>>
485        + Shr<u32, Output = Saturating<Self>>
486        + AddAssign
487        + SubAssign
488        + MulAssign
489        + DivAssign
490        + RemAssign
491        + BitAndAssign
492        + BitOrAssign
493        + BitXorAssign
494        + ShlAssign<u32>
495        + ShrAssign<u32>
496        + Not<Output = Saturating<Self>>
497        + Add<Self, Output = Saturating<Self>>
498        + Sub<Self, Output = Saturating<Self>>
499        + Mul<Self, Output = Saturating<Self>>
500        + Div<Self, Output = Saturating<Self>>
501        + Rem<Self, Output = Saturating<Self>>
502        + BitAnd<Self, Output = Saturating<Self>>
503        + BitOr<Self, Output = Saturating<Self>>
504        + BitXor<Self, Output = Saturating<Self>>
505        + AddAssign<Self>
506        + SubAssign<Self>
507        + MulAssign<Self>
508        + DivAssign<Self>
509        + RemAssign<Self>
510        + BitAndAssign<Self>
511        + BitOrAssign<Self>
512        + BitXorAssign<Self>,
513{
514    fn to_saturating(self) -> Saturating<Self> {
515        Saturating(self)
516    }
517    fn from_saturating(s: Saturating<Self>) -> Self {
518        s.0
519    }
520}
521
522impl<T> fmt::Debug for Saturating<T>
523where
524    T: fmt::Debug,
525{
526    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
527        T::fmt(&self.0, f)
528    }
529}
530impl<T> Bounded for Saturating<T>
531where
532    T: Bounded,
533{
534    fn maximum() -> Self {
535        Self(T::maximum())
536    }
537    fn minimum() -> Self {
538        Self(T::minimum())
539    }
540}
541impl<T> Zero for Saturating<T>
542where
543    T: Zero,
544{
545    fn zero() -> Self {
546        Self(T::zero())
547    }
548}
549impl<T> One for Saturating<T>
550where
551    T: One,
552{
553    fn one() -> Self {
554        Self(T::one())
555    }
556}
557impl<T> FromStr for Saturating<T>
558where
559    T: FromStr,
560{
561    type Err = T::Err;
562    fn from_str(s: &str) -> Result<Self, Self::Err> {
563        T::from_str(s).map(Self)
564    }
565}
566impl<T> Display for Saturating<T>
567where
568    T: Display,
569{
570    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
571        T::fmt(&self.0, f)
572    }
573}
574impl<T> IterScan for Saturating<T>
575where
576    T: IterScan<Output = T>,
577{
578    type Output = Self;
579    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
580        T::scan(iter).map(Self)
581    }
582}
583impl_binop!(impl<T> Div div for Saturating<T>);
584impl_binop!(impl<T> Rem rem for Saturating<T>);
585impl_binop!(impl<T> BitAnd bitand for Saturating<T>);
586impl_binop!(impl<T> BitOr bitor for Saturating<T>);
587impl_binop!(impl<T> BitXor bitxor for Saturating<T>);
588impl_opassign!(impl<T> AddAssign add_assign for Saturating<T> => Add add);
589impl_opassign!(impl<T> SubAssign sub_assign for Saturating<T> => Sub sub);
590impl_opassign!(impl<T> MulAssign mul_assign for Saturating<T> => Mul mul);
591impl_opassign!(impl<T> DivAssign div_assign for Saturating<T>);
592impl_opassign!(impl<T> RemAssign rem_assign for Saturating<T>);
593impl_opassign!(impl<T> BitAndAssign bitand_assign for Saturating<T>);
594impl_opassign!(impl<T> BitOrAssign bitor_assign for Saturating<T>);
595impl_opassign!(impl<T> BitXorAssign bitxor_assign for Saturating<T>);
596impl<T> Not for Saturating<T>
597where
598    T: Not<Output = T>,
599{
600    type Output = Self;
601    fn not(self) -> Self::Output {
602        Self(Not::not(self.0))
603    }
604}
605
606macro_rules! impl_int_base_for_saturating {
607    ($($t:ty)*) => {
608        $(
609            impl Saturatingable for $t {}
610            impl Add for Saturating<$t> {
611                type Output = Self;
612                fn add(self, rhs: Self) -> Self::Output {
613                    Self(self.0.saturating_add(rhs.0))
614                }
615            }
616            impl Add<$t> for Saturating<$t> {
617                type Output = Self;
618                fn add(self, rhs: $t) -> Self::Output {
619                    Self(self.0.saturating_add(rhs))
620                }
621            }
622            impl Sub for Saturating<$t> {
623                type Output = Self;
624                fn sub(self, rhs: Self) -> Self::Output {
625                    Self(self.0.saturating_sub(rhs.0))
626                }
627            }
628            impl Sub<$t> for Saturating<$t> {
629                type Output = Self;
630                fn sub(self, rhs: $t) -> Self::Output {
631                    Self(self.0.saturating_sub(rhs))
632                }
633            }
634            impl Mul for Saturating<$t> {
635                type Output = Self;
636                fn mul(self, rhs: Self) -> Self::Output {
637                    Self(self.0.saturating_mul(rhs.0))
638                }
639            }
640            impl Mul<$t> for Saturating<$t> {
641                type Output = Self;
642                fn mul(self, rhs: $t) -> Self::Output {
643                    Self(self.0.saturating_mul(rhs))
644                }
645            }
646            impl Sum for Saturating<$t> {
647                fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
648                    iter.fold(Self::zero(), |acc, x| acc + x)
649                }
650            }
651            impl Product for Saturating<$t> {
652                fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
653                    iter.fold(Self::one(), |acc, x| acc * x)
654                }
655            }
656            impl IntBase for Saturating<$t> {
657                type Error = <$t as IntBase>::Error;
658                fn div_euclid(self, rhs: Self) -> Self { Self(self.0.div_euclid(rhs.0)) }
659                fn rem_euclid(self, rhs: Self) -> Self { Self(self.0.rem_euclid(rhs.0)) }
660                fn pow(self, exp: u32) -> Self { Self(self.0.saturating_pow(exp)) }
661                fn from_str_radix(src: &str, radix: u32) -> Result<Self, Self::Error> { <$t as IntBase>::from_str_radix(src, radix).map(Self) }
662                fn ilog(self, base: Self) -> u32 { self.0.ilog(base.0) }
663                fn ilog2(self) -> u32 { self.0.ilog2() }
664                fn ilog10(self) -> u32 { self.0.ilog10() }
665            }
666            impl From<$t> for Saturating<$t> {
667                fn from(t: $t) -> Self {
668                    Self(t)
669                }
670            }
671        )*
672    };
673}
674impl_int_base_for_saturating!(u8 i8 u16 i16 u32 i32 u64 i64 u128 i128 usize isize);
675
676macro_rules! impl_unsigned_signed_for_saturating {
677    ($($unsigned:ident $signed:ident)*) => {
678        $(
679            impl Unsigned for Saturating<$unsigned> {
680                type Signed = Saturating<$signed>;
681                fn signed(self) -> Self::Signed { Saturating(TryFrom::try_from(self.0).ok().unwrap_or_else($signed::maximum)) }
682                fn abs_diff(self, other: Self) -> Self { Self(self.0.abs_diff(other.0)) }
683                fn next_power_of_two(self) -> Self { Self(self.0.next_power_of_two()) }
684                fn gcd(self, other: Self) -> Self { Self(self.0.gcd(other.0)) }
685                fn mod_add(self, rhs: Self, modulo: Self) -> Self { Self(self.0.mod_add(rhs.0, modulo.0)) }
686                fn mod_sub(self, rhs: Self, modulo: Self) -> Self { Self(self.0.mod_sub(rhs.0, modulo.0)) }
687                fn mod_mul(self, rhs: Self, modulo: Self) -> Self { Self(self.0.mod_mul(rhs.0, modulo.0)) }
688            }
689            impl Signed for Saturating<$signed> {
690                type Unsigned = Saturating<$unsigned>;
691                fn unsigned(self) -> Self::Unsigned { Saturating(TryFrom::try_from(self.0).ok().unwrap_or_else($unsigned::minimum)) }
692                fn abs(self) -> Self { Self(self.0.saturating_abs()) }
693                fn abs_diff(self, other: Self) -> Self::Unsigned { Saturating(self.0.abs_diff(other.0)) }
694                fn is_negative(self) -> bool { self.0.is_negative() }
695                fn is_positive(self) -> bool { self.0.is_positive() }
696                fn signum(self) -> Self { Self(self.0.signum()) }
697            }
698            impl Neg for Saturating<$signed> {
699                type Output = Self;
700                fn neg(self) -> Self::Output {
701                    Self(self.0.saturating_neg())
702                }
703            }
704        )*
705    };
706}
707impl_unsigned_signed_for_saturating!(u8 i8 u16 i16 u32 i32 u64 i64 u128 i128 usize isize);
708
709macro_rules! impl_binary_repr_for_saturating {
710    ($($t:ty)*) => {
711        $(
712            impl Shl<u32> for Saturating<$t> {
713                type Output = Self;
714                fn shl(self, rhs: u32) -> Self::Output {
715                    Self(self.0.checked_shl(rhs).unwrap_or(0))
716                }
717            }
718            impl Shr<u32> for Saturating<$t> {
719                type Output = Self;
720                fn shr(self, rhs: u32) -> Self::Output {
721                    Self(self.0.checked_shr(rhs).unwrap_or(0))
722                }
723            }
724            impl ShlAssign<u32> for Saturating<$t> {
725                fn shl_assign(&mut self, rhs: u32) {
726                    *self = Shl::shl(*self, rhs);
727                }
728            }
729            impl ShrAssign<u32> for Saturating<$t> {
730                fn shr_assign(&mut self, rhs: u32) {
731                    *self = Shr::shr(*self, rhs);
732                }
733            }
734            impl BinaryRepr for Saturating<$t> {
735                fn count_ones(self) -> u32 { self.0.count_ones() }
736                fn count_zeros(self) -> u32 { self.0.count_zeros() }
737                fn leading_ones(self) -> u32 { self.0.leading_ones() }
738                fn leading_zeros(self) -> u32 { self.0.leading_zeros() }
739                fn reverse_bits(self) -> Self { Self(self.0.reverse_bits()) }
740                fn rotate_left(self, n: u32) -> Self { Self(self.0.rotate_left(n)) }
741                fn rotate_right(self, n: u32) -> Self { Self(self.0.rotate_right(n)) }
742                fn swap_bytes(self) -> Self { Self(self.0.swap_bytes()) }
743                fn trailing_ones(self) -> u32 { self.0.trailing_ones() }
744                fn trailing_zeros(self) -> u32 { self.0.trailing_zeros() }
745            }
746        )*
747    };
748}
749impl_binary_repr_for_saturating!(u8 i8 u16 i16 u32 i32 u64 i64 u128 i128 usize isize);
750
751#[derive(Default, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
752#[repr(transparent)]
753/// Wrapper type of arithmetic `wrapping_*` operations.
754pub struct Wrapping<T>(pub T);
755pub trait Wrappingable: Sized
756where
757    Wrapping<Self>: Copy
758        + Bounded
759        + Zero
760        + One
761        + Eq
762        + Ord
763        + Default
764        + FromStr
765        + Display
766        + Add<Output = Wrapping<Self>>
767        + Sub<Output = Wrapping<Self>>
768        + Mul<Output = Wrapping<Self>>
769        + Div<Output = Wrapping<Self>>
770        + Rem<Output = Wrapping<Self>>
771        + BitAnd<Output = Wrapping<Self>>
772        + BitOr<Output = Wrapping<Self>>
773        + BitXor<Output = Wrapping<Self>>
774        + Shl<u32, Output = Wrapping<Self>>
775        + Shr<u32, Output = Wrapping<Self>>
776        + AddAssign
777        + SubAssign
778        + MulAssign
779        + DivAssign
780        + RemAssign
781        + BitAndAssign
782        + BitOrAssign
783        + BitXorAssign
784        + ShlAssign<u32>
785        + ShrAssign<u32>
786        + Not<Output = Wrapping<Self>>
787        + Add<Self, Output = Wrapping<Self>>
788        + Sub<Self, Output = Wrapping<Self>>
789        + Mul<Self, Output = Wrapping<Self>>
790        + Div<Self, Output = Wrapping<Self>>
791        + Rem<Self, Output = Wrapping<Self>>
792        + BitAnd<Self, Output = Wrapping<Self>>
793        + BitOr<Self, Output = Wrapping<Self>>
794        + BitXor<Self, Output = Wrapping<Self>>
795        + AddAssign<Self>
796        + SubAssign<Self>
797        + MulAssign<Self>
798        + DivAssign<Self>
799        + RemAssign<Self>
800        + BitAndAssign<Self>
801        + BitOrAssign<Self>
802        + BitXorAssign<Self>,
803{
804    fn to_wrapping(self) -> Wrapping<Self> {
805        Wrapping(self)
806    }
807    fn from_wrapping(w: Wrapping<Self>) -> Self {
808        w.0
809    }
810}
811
812impl<T> fmt::Debug for Wrapping<T>
813where
814    T: fmt::Debug,
815{
816    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
817        T::fmt(&self.0, f)
818    }
819}
820impl<T> Bounded for Wrapping<T>
821where
822    T: Bounded,
823{
824    fn maximum() -> Self {
825        Self(T::maximum())
826    }
827    fn minimum() -> Self {
828        Self(T::minimum())
829    }
830}
831impl<T> Zero for Wrapping<T>
832where
833    T: Zero,
834{
835    fn zero() -> Self {
836        Self(T::zero())
837    }
838}
839impl<T> One for Wrapping<T>
840where
841    T: One,
842{
843    fn one() -> Self {
844        Self(T::one())
845    }
846}
847impl<T> FromStr for Wrapping<T>
848where
849    T: FromStr,
850{
851    type Err = T::Err;
852    fn from_str(s: &str) -> Result<Self, Self::Err> {
853        T::from_str(s).map(Self)
854    }
855}
856impl<T> Display for Wrapping<T>
857where
858    T: Display,
859{
860    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
861        T::fmt(&self.0, f)
862    }
863}
864impl<T> IterScan for Wrapping<T>
865where
866    T: IterScan<Output = T>,
867{
868    type Output = Self;
869    fn scan<'a, I: Iterator<Item = &'a str>>(iter: &mut I) -> Option<Self::Output> {
870        T::scan(iter).map(Self)
871    }
872}
873impl_binop!(impl<T> BitAnd bitand for Wrapping<T>);
874impl_binop!(impl<T> BitOr bitor for Wrapping<T>);
875impl_binop!(impl<T> BitXor bitxor for Wrapping<T>);
876impl_opassign!(impl<T> AddAssign add_assign for Wrapping<T> => Add add);
877impl_opassign!(impl<T> SubAssign sub_assign for Wrapping<T> => Sub sub);
878impl_opassign!(impl<T> MulAssign mul_assign for Wrapping<T> => Mul mul);
879impl_opassign!(impl<T> DivAssign div_assign for Wrapping<T> => Div div);
880impl_opassign!(impl<T> RemAssign rem_assign for Wrapping<T> => Rem rem);
881impl_opassign!(impl<T> BitAndAssign bitand_assign for Wrapping<T>);
882impl_opassign!(impl<T> BitOrAssign bitor_assign for Wrapping<T>);
883impl_opassign!(impl<T> BitXorAssign bitxor_assign for Wrapping<T>);
884impl<T> Not for Wrapping<T>
885where
886    T: Not<Output = T>,
887{
888    type Output = Self;
889    fn not(self) -> Self::Output {
890        Self(Not::not(self.0))
891    }
892}
893
894macro_rules! impl_int_base_for_wrapping {
895    ($($t:ty)*) => {
896        $(
897            impl Wrappingable for $t {}
898            impl Add for Wrapping<$t> {
899                type Output = Self;
900                fn add(self, rhs: Self) -> Self::Output {
901                    Self(self.0.wrapping_add(rhs.0))
902                }
903            }
904            impl Add<$t> for Wrapping<$t> {
905                type Output = Self;
906                fn add(self, rhs: $t) -> Self::Output {
907                    Self(self.0.wrapping_add(rhs))
908                }
909            }
910            impl Sub for Wrapping<$t> {
911                type Output = Self;
912                fn sub(self, rhs: Self) -> Self::Output {
913                    Self(self.0.wrapping_sub(rhs.0))
914                }
915            }
916            impl Sub<$t> for Wrapping<$t> {
917                type Output = Self;
918                fn sub(self, rhs: $t) -> Self::Output {
919                    Self(self.0.wrapping_sub(rhs))
920                }
921            }
922            impl Mul for Wrapping<$t> {
923                type Output = Self;
924                fn mul(self, rhs: Self) -> Self::Output {
925                    Self(self.0.wrapping_mul(rhs.0))
926                }
927            }
928            impl Mul<$t> for Wrapping<$t> {
929                type Output = Self;
930                fn mul(self, rhs: $t) -> Self::Output {
931                    Self(self.0.wrapping_mul(rhs))
932                }
933            }
934            impl Div for Wrapping<$t> {
935                type Output = Self;
936                fn div(self, rhs: Self) -> Self::Output {
937                    Self(self.0.wrapping_div(rhs.0))
938                }
939            }
940            impl Div<$t> for Wrapping<$t> {
941                type Output = Self;
942                fn div(self, rhs: $t) -> Self::Output {
943                    Self(self.0.wrapping_div(rhs))
944                }
945            }
946            impl Rem for Wrapping<$t> {
947                type Output = Self;
948                fn rem(self, rhs: Self) -> Self::Output {
949                    Self(self.0.wrapping_rem(rhs.0))
950                }
951            }
952            impl Rem<$t> for Wrapping<$t> {
953                type Output = Self;
954                fn rem(self, rhs: $t) -> Self::Output {
955                    Self(self.0.wrapping_rem(rhs))
956                }
957            }
958            impl Sum for Wrapping<$t> {
959                fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
960                    iter.fold(Self::zero(), |acc, x| acc + x)
961                }
962            }
963            impl Product for Wrapping<$t> {
964                fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
965                    iter.fold(Self::one(), |acc, x| acc * x)
966                }
967            }
968            impl IntBase for Wrapping<$t> {
969                type Error = <$t as IntBase>::Error;
970                fn div_euclid(self, rhs: Self) -> Self { Self(self.0.wrapping_div_euclid(rhs.0)) }
971                fn rem_euclid(self, rhs: Self) -> Self { Self(self.0.wrapping_rem_euclid(rhs.0)) }
972                fn pow(self, exp: u32) -> Self { Self(self.0.wrapping_pow(exp)) }
973                fn from_str_radix(src: &str, radix: u32) -> Result<Self, Self::Error> { <$t as IntBase>::from_str_radix(src, radix).map(Self) }
974                fn ilog(self, base: Self) -> u32 { self.0.ilog(base.0) }
975                fn ilog2(self) -> u32 { self.0.ilog2() }
976                fn ilog10(self) -> u32 { self.0.ilog10() }
977            }
978            impl From<$t> for Wrapping<$t> {
979                fn from(t: $t) -> Self {
980                    Self(t)
981                }
982            }
983        )*
984    };
985}
986impl_int_base_for_wrapping!(u8 i8 u16 i16 u32 i32 u64 i64 u128 i128 usize isize);
987
988macro_rules! impl_unsigned_signed_for_wrapping {
989    ($($unsigned:ident $signed:ident)*) => {
990        $(
991            impl Unsigned for Wrapping<$unsigned> {
992                type Signed = Wrapping<$signed>;
993                fn signed(self) -> Self::Signed { Wrapping(self.0.signed()) }
994                fn abs_diff(self, other: Self) -> Self { Self(self.0.abs_diff(other.0)) }
995                fn next_power_of_two(self) -> Self { Self(self.0.next_power_of_two()) }
996                fn gcd(self, other: Self) -> Self { Self(self.0.gcd(other.0)) }
997                fn mod_add(self, rhs: Self, modulo: Self) -> Self { Self(self.0.mod_add(rhs.0, modulo.0)) }
998                fn mod_sub(self, rhs: Self, modulo: Self) -> Self { Self(self.0.mod_sub(rhs.0, modulo.0)) }
999                fn mod_mul(self, rhs: Self, modulo: Self) -> Self { Self(self.0.mod_mul(rhs.0, modulo.0)) }
1000            }
1001            impl Signed for Wrapping<$signed> {
1002                type Unsigned = Wrapping<$unsigned>;
1003                fn unsigned(self) -> Self::Unsigned { Wrapping(self.0.unsigned()) }
1004                fn abs(self) -> Self { Self(self.0.wrapping_abs()) }
1005                fn abs_diff(self, other: Self) -> Self::Unsigned { Wrapping(self.0.abs_diff(other.0)) }
1006                fn is_negative(self) -> bool { self.0.is_negative() }
1007                fn is_positive(self) -> bool { self.0.is_positive() }
1008                fn signum(self) -> Self { Self(self.0.signum()) }
1009            }
1010            impl Neg for Wrapping<$signed> {
1011                type Output = Self;
1012                fn neg(self) -> Self::Output {
1013                    Self(self.0.wrapping_neg())
1014                }
1015            }
1016        )*
1017    };
1018}
1019impl_unsigned_signed_for_wrapping!(u8 i8 u16 i16 u32 i32 u64 i64 u128 i128 usize isize);
1020
1021macro_rules! impl_binary_repr_for_wrapping {
1022    ($($t:ty)*) => {
1023        $(
1024            impl Shl<u32> for Wrapping<$t> {
1025                type Output = Self;
1026                fn shl(self, rhs: u32) -> Self::Output {
1027                    Self(self.0.wrapping_shl(rhs))
1028                }
1029            }
1030            impl Shr<u32> for Wrapping<$t> {
1031                type Output = Self;
1032                fn shr(self, rhs: u32) -> Self::Output {
1033                    Self(self.0.wrapping_shr(rhs))
1034                }
1035            }
1036            impl ShlAssign<u32> for Wrapping<$t> {
1037                fn shl_assign(&mut self, rhs: u32) {
1038                    *self = Shl::shl(*self, rhs);
1039                }
1040            }
1041            impl ShrAssign<u32> for Wrapping<$t> {
1042                fn shr_assign(&mut self, rhs: u32) {
1043                    *self = Shr::shr(*self, rhs);
1044                }
1045            }
1046            impl BinaryRepr for Wrapping<$t> {
1047                fn count_ones(self) -> u32 { self.0.count_ones() }
1048                fn count_zeros(self) -> u32 { self.0.count_zeros() }
1049                fn leading_ones(self) -> u32 { self.0.leading_ones() }
1050                fn leading_zeros(self) -> u32 { self.0.leading_zeros() }
1051                fn reverse_bits(self) -> Self { Self(self.0.reverse_bits()) }
1052                fn rotate_left(self, n: u32) -> Self { Self(self.0.rotate_left(n)) }
1053                fn rotate_right(self, n: u32) -> Self { Self(self.0.rotate_right(n)) }
1054                fn swap_bytes(self) -> Self { Self(self.0.swap_bytes()) }
1055                fn trailing_ones(self) -> u32 { self.0.trailing_ones() }
1056                fn trailing_zeros(self) -> u32 { self.0.trailing_zeros() }
1057            }
1058        )*
1059    };
1060}
1061impl_binary_repr_for_wrapping!(u8 i8 u16 i16 u32 i32 u64 i64 u128 i128 usize isize);
1062
1063#[cfg(test)]
1064mod tests {
1065    use super::*;
1066    use crate::tools::Xorshift;
1067    const Q: usize = 10_000;
1068    macro_rules! test_unsigned {
1069        ($($t:ident)*) => {
1070            $(
1071                mod $t {
1072                    use super::*;
1073                    const A: $t = $t::MAX / 2;
1074                    fn gcd(mut a: $t, mut b: $t) -> $t {
1075                        while b != 0 {
1076                            a %= b;
1077                            std::mem::swap(&mut a, &mut b);
1078                        }
1079                        a
1080                    }
1081                    #[test]
1082                    fn test_gcd() {
1083                        let mut rng = Xorshift::default();
1084                            for (a, b) in rng.random_iter((0..=A, 0..=A)).take(Q) {
1085                            assert_eq!(a.gcd(b), gcd(a, b));
1086                        }
1087                        assert_eq!($t::zero().gcd(0), 0);
1088                        assert_eq!($t::zero().gcd(100), 100);
1089                    }
1090                    #[test]
1091                    fn test_mod_inv() {
1092                        let mut rng = Xorshift::default();
1093                        for _ in 0..Q {
1094                            let m = rng.random(2..=A);
1095                            let a = rng.random(1..m);
1096                            let g = a.gcd(m);
1097                            let m = m / g;
1098                            let a = a / g;
1099                            let x = a.mod_inv(m);
1100                            assert!(x < m);
1101                            assert_eq!(a as u128 * x as u128 % m as u128, 1);
1102                        }
1103                    }
1104                    #[test]
1105                    fn test_mod_operate() {
1106                        let mut rng = Xorshift::default();
1107                        for _ in 0..Q {
1108                            for ub in [10, A] {
1109                                let m = rng.random(2..=ub);
1110                                let a = rng.random(0..m);
1111                                let b = rng.random(0..m);
1112                                assert_eq!(a.mod_add(b, m), ((a as u128 + b as u128) % m as u128) as $t);
1113                                assert_eq!(a.mod_sub(b, m), ((a as u128 + m as u128 - b as u128) % m as u128) as $t);
1114                                assert_eq!(a.mod_mul(b, m), (a as u128 * b as u128 % m as u128) as $t);
1115                                assert_eq!(a.mod_mul(b, m), (a as u128).mod_mul(b as u128, m as u128) as $t);
1116                                assert_eq!(a.mod_neg(m), ((m as u128 - a as u128) % m as u128) as $t);
1117                            }
1118                        }
1119                    }
1120                }
1121            )*
1122        };
1123    }
1124    test_unsigned!(u8 u16 u32 u64 usize);
1125
1126    macro_rules! test_signed {
1127        ($($t:ident)*) => {
1128            $(
1129                mod $t {
1130                    use super::*;
1131                    const A: $t = $t::MAX / 2;
1132                    #[test]
1133                    fn test_extgcd() {
1134                        let mut rng = Xorshift::default();
1135                        for (a, b) in rng.random_iter((-A..=A, -A..=A)).take(Q) {
1136                            let ExtendedGcd { g, x, y } = a.extgcd(b);
1137                            assert_eq!(g, a.abs().unsigned().gcd(b.abs().unsigned()));
1138                            assert_eq!(a as i128 * x as i128 + b as i128 * y as i128, g.signed() as i128);
1139                        }
1140                    }
1141                }
1142            )*
1143        };
1144    }
1145    test_signed!(i8 i16 i32 i64 isize);
1146
1147    #[test]
1148    fn test_mod_mul_u128() {
1149        fn naive_mod_mul(a: u128, b: u128, m: u128) -> u128 {
1150            assert!(m != 0);
1151            let a = [a as u64, (a >> 64) as u64];
1152            let b = [b as u64, (b >> 64) as u64];
1153            let mut res = 0u128;
1154            for (i, &a) in a.iter().enumerate() {
1155                for (j, &b) in b.iter().enumerate() {
1156                    let mut x = (a as u128) * (b as u128) % m;
1157                    for _ in 0..(i + j) * 64 {
1158                        x = x.mod_add(x, m);
1159                    }
1160                    res = res.mod_add(x, m);
1161                }
1162            }
1163            res
1164        }
1165
1166        let mut rng = Xorshift::default();
1167        for _ in 0..Q {
1168            for a in [1, 10, u32::MAX as _, u64::MAX as _, u128::MAX] {
1169                for b in [1, 10, u32::MAX as _, u64::MAX as _, u128::MAX] {
1170                    for c in [1, 10, u32::MAX as _, u64::MAX as _, u128::MAX] {
1171                        let m = rng.random(1..=c);
1172                        let x = rng.random(0..a.min(m));
1173                        let y = rng.random(0..b.min(m));
1174                        assert_eq!(x.mod_mul(y, m), naive_mod_mul(x, y, m));
1175                        let x = rng.random(0..a);
1176                        let y = rng.random(0..b);
1177                        assert_eq!(x.mod_mul(y, m), naive_mod_mul(x, y, m));
1178                    }
1179                }
1180            }
1181        }
1182    }
1183
1184    #[test]
1185    fn test_rem() {
1186        fn naive_rem(u: [u64; 4], v: u128) -> u128 {
1187            assert!(v != 0);
1188            let mut u = [
1189                ((u[1] as u128) << 64) | (u[0] as u128),
1190                ((u[3] as u128) << 64) | (u[2] as u128),
1191            ];
1192            let mut v_mul_2 = vec![[v, 0]];
1193            while v_mul_2.last().unwrap()[1].leading_zeros() != 0 {
1194                let [v_lo, v_hi] = *v_mul_2.last().unwrap();
1195                v_mul_2.push([v_lo << 1, v_hi << 1 | (v_lo >> 127)]);
1196            }
1197            v_mul_2.reverse();
1198            for [v_lo, v_hi] in v_mul_2 {
1199                let [u_lo, u_hi] = u;
1200                if (u_hi > v_hi) || (u_hi == v_hi && u_lo >= v_lo) {
1201                    let (new_lo, carry) = u_lo.overflowing_sub(v_lo);
1202                    let new_hi = u_hi - v_hi - (carry as u128);
1203                    u = [new_lo, new_hi];
1204                }
1205            }
1206            u[0]
1207        }
1208        let mut rng = Xorshift::default();
1209        for _ in 0..Q {
1210            let mut u = [0u64; 4];
1211            for k in 0..4 {
1212                for a in [1, 10, u32::MAX as _, u64::MAX] {
1213                    u[k] = rng.random(..a);
1214                    for b in [1, 10, u128::MAX] {
1215                        let v = rng.random(1..=b);
1216                        assert_eq!(rem_u256_by_u128(u, v), naive_rem(u, v));
1217                    }
1218                }
1219            }
1220        }
1221    }
1222}
competitive/num/integer.rs

competitive/num/
integer.rs
