atcoder_zp.rs

mod zp{
    use super::arith_util;
    
    #[derive(Debug, Clone, Copy, Eq, PartialEq)]
    pub struct Zp{
        pub val: i64,
        pub p: i64
    }

    impl Zp{
        pub fn new(val: i64, p: i64) -> Self{
            assert!(p > 0);
            let (_, r) = arith_util::div(val, p);
            return Zp{val: r, p: p};
        }

        pub fn inv_safe(&self) -> Option<Self>{
            if self.val == 0{
                return None;
            }
            // rhsの逆数を
            let ((x, y), g) = arith_util::extgcd(self.val, self.p);
            if g != 1{
                return None;
            }

            return Some(Self::new(x, self.p));
        }

        pub fn div_safe(&self, rhs: Self) -> Option<Self>{
            if self.p != rhs.p{
                return None;
            }
            return Some(*self*rhs.inv_safe()?);
        }
    }

    impl std::ops::Add for Zp{
        type Output = Self;

        fn add(self, rhs: Self) -> Self{
            assert_eq!(self.p, rhs.p);
            return Zp::new(self.val + rhs.val, self.p);
        }
    }

    impl std::ops::AddAssign for Zp{
        fn add_assign(&mut self, rhs: Self) {
            *self = *self + rhs;
        }
    }

    impl std::ops::Neg for Zp {
        type Output = Self;
    
        fn neg(self) -> Self {
            return Zp::new(-self.val, self.p);
        }
    }

    impl std::ops::Sub for Zp {
        type Output = Self;

        fn sub(self, other: Self) -> Self::Output {
            return self + (-other);
        }
    }

    impl std::ops::SubAssign for Zp {
        fn sub_assign(&mut self, rhs: Self) {
            *self = self.clone() - rhs;
        }
    }

    impl std::ops::Mul for Zp{
        type Output = Self;

        fn mul(self, rhs: Self) -> Self::Output {
            assert_eq!(self.p, rhs.p);
            return Zp::new(self.val * rhs.val, self.p);
        }
    }

    impl std::ops::MulAssign for Zp {
        fn mul_assign(&mut self, rhs: Self) {
            *self = self.clone() * rhs;
        }
    }

    impl std::ops::Div for Zp{
        type Output = Self;

        fn div(self, rhs: Self) -> Self::Output{
            let res = self.div_safe(rhs);
            return res.unwrap();
        }
    }

    impl std::ops::DivAssign for Zp{
        fn div_assign(&mut self, rhs: Self){
            *self = self.clone() / rhs;
        }
    }

    impl std::fmt::Display for Zp {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(f, "[{}]{}", self.val, self.p)
        }
    }

    #[test]
    fn test(){
        assert_eq!(Zp::new(3, 5) + Zp::new(4, 5), Zp::new(2, 5));
        assert_eq!(-Zp::new(3, 5), Zp::new(2, 5));
        assert_eq!(Zp::new(1, 5) - Zp::new(3, 5), Zp::new(3, 5));
        assert_eq!(Zp::new(4, 5) * Zp::new(4, 5), Zp::new(1, 5));
        assert_eq!(Zp::new(1, 5) / Zp::new(4, 5), Zp::new(4, 5));
        assert_eq!(Zp::new(1, 5) / Zp::new(3, 5), Zp::new(2, 5));
    }

}
use zp::Zp;