#799286 (Rust) No.3030 ミラー・ラビン素数判定法のテスト

• 問題: No.3030 ミラー・ラビン素数判定法のテスト
• ユーザー: 👑 Mizar
• 提出日時: 2022-09-20 11:46:21
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 28 ms / 9,973 ms
• コード長: 7,197 bytes
• コンパイル時間: 12,398 ms
• コンパイル使用メモリ: 401,732 KB
• 実行使用メモリ: 6,944 KB
• 最終ジャッジ日時: 2024-06-01 16:56:37
• 合計ジャッジ時間: 13,386 ms
• ジャッジサーバーID （参考情報）: judge2 / judge1

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 6,812 KB
testcase_01	AC	1 ms; 6,816 KB
testcase_02	AC	1 ms; 6,940 KB
testcase_03	AC	1 ms; 6,940 KB
testcase_04	AC	17 ms; 6,944 KB
testcase_05	AC	17 ms; 6,940 KB
testcase_06	AC	9 ms; 6,940 KB
testcase_07	AC	9 ms; 6,940 KB
testcase_08	AC	9 ms; 6,944 KB
testcase_09	AC	28 ms; 6,944 KB

ソースコード

// -*- coding:utf-8-unix -*-

// モンゴメリ剰余乗算 (Montgomery modular multiplication)
pub struct Mont<T, BitCounts=u32> {
    n: T, // n is odd, and n > 2
    ni: T, // n * ni == 1 (mod 2**64 {@T=u64})
    r: T, // == 2**64 {@T=u64} (mod n)
    rn: T, // == -1 * (2**64 {T=u64}) (mod n)
    r2: T, // == 2**128 {@T=u64} (mod n)
    d: T, // == (n - 1) >> (n - 1).trailing_zeros()
    k: BitCounts, // == (n - 1).trailing_zeros()
}
pub trait MontTrait<T, BitCounts=u32> {
    fn n(&self) -> T;
    fn ni(&self) -> T;
    fn r(&self) -> T;
    fn rn(&self) -> T;
    fn r2(&self) -> T;
    fn d(&self) -> T;
    fn k(&self) -> BitCounts;
    // create constant structures for montgomery modular arithmetic
    fn new(n: T) -> Self;
    // addmod(a, b) == a + b (mod n)
    fn add(&self, a: T, b: T) -> T;
    // submod(a, b) == a - b (mod n)
    fn sub(&self, a: T, b: T) -> T;
    // mrmul(ar, br) == (ar * br) / r (mod n)
    // R == 2**64 {@T=u64}
    // gcd(N, R) == 1
    // N * ni mod R == 1
    // 0 <= ar < N < R
    // 0 <= br < N < R
    // T := ar * br
    // t := floor(T / R) - floor(((T * ni mod R) * N) / R)
    // if t < 0 then return t + N else return t
    fn mrmul(&self, ar: T, br: T) -> T;
    // mr(ar) == ar / r (mod n)
    // R == 2**64 {@T=u64}
    // gcd(N, R) == 1
    // N * ni mod R == 1
    // 0 <= ar < N < R
    // t := floor(ar / R) - floor(((ar * ni mod R) * N) / R)
    // if t < 0 then return t + N else return t
    fn mr(&self, ar: T) -> T;
    // ar(a) == a * r (mod n)
    fn ar(&self, a: T) -> T;
    // powir(ar, b) == ((ar / r) ** b) * r (mod n)
    fn pow(&self, ar: T, b: T) -> T;
}
impl MontTrait<u64> for Mont<u64> {
    #[inline] fn n(&self) -> u64 { self.n }
    #[inline] fn ni(&self) -> u64 { self.ni }
    #[inline] fn r(&self) -> u64 { self.r }
    #[inline] fn rn(&self) -> u64 { self.rn }
    #[inline] fn r2(&self) -> u64 { self.r2 }
    #[inline] fn d(&self) -> u64 { self.d }
    #[inline] fn k(&self) -> u32 { self.k }
    #[inline]
    fn new(n: u64) -> Self {
        // create constant structures for montgomery modular arithmetic
        debug_assert_eq!(n & 1, 1);
        // // n is odd number, n = 2*k+1, n >= 1, n < 2**64, k is non-negative integer, k >= 0, k < 2**63
        // ni0 := n; // = 2*k+1 = (1+(2**3)*((k*(k+1)/2)**1))/(2*k+1)
        let mut ni = n;
        // ni1 := ni0 * ((n * ni0 - 3) * (n * ni0) + 3); // = (1+(2**9)*((k*(k+1)/2)**3))/(2*k+1)
        // ni2 := ni1 * ((n * ni1 - 3) * (n * ni1) + 3); // = (1+(2**27)*((k*(k+1)/2)**9))/(2*k+1)
        // ni3 := ni2 * ((n * ni2 - 3) * (n * ni2) + 3); // = (1+(2**81)*((k*(k+1)/2)**27))/(2*k+1)
        // // (n * ni3) mod 2**64 = ((2*k+1) * ni3) mod 2**64 = 1 mod 2**64
        for _ in 0..3 {
            let t = n.wrapping_mul(ni);
            ni = t.wrapping_sub(3).wrapping_mul(t).wrapping_add(3).wrapping_mul(ni);
        }
        debug_assert_eq!(n.wrapping_mul(ni), 1); // n * ni == 1 (mod 2**64)
        let r: u64 = n.wrapping_neg() % n; // == 2**64 (mod n)
        let rn: u64 = n - r; // == -1 * (2**64) (mod n)
        let r2: u64 = ((n as u128).wrapping_neg() % (n as u128)) as u64; // == 2**128 (mod n)
        // n == 2**k * d + 1
        let mut d = n - 1;
        let k = d.trailing_zeros();
        d >>= k;
        debug_assert_eq!(Self { n, ni, r, rn, r2, d, k }.mr(r), 1); // r / r == 1 (mod n)
        debug_assert_eq!(Self { n, ni, r, rn, r2, d, k }.mrmul(1, r2), r); // r2 / r == r (mod n)
        Self { n, ni, r, rn, r2, d, k }
    }
    #[inline]
    fn add(&self, a: u64, b: u64) -> u64 {
        // addmod(a, b) == a + b (mod n)
        debug_assert!(a < self.n());
        debug_assert!(b < self.n());
        let (t, fa) = a.overflowing_add(b);
        let (u, fs) = t.overflowing_sub(self.n());
        if fa || !fs { u } else { t }
    }
    #[inline]
    fn sub(&self, a: u64, b: u64) -> u64 {
        // submod(a, b) == a - b (mod n)
        debug_assert!(a < self.n());
        debug_assert!(b < self.n());
        let (t, f) = a.overflowing_sub(b);
        if f { t.wrapping_add(self.n()) } else { t }
    }
    #[inline]
    fn mrmul(&self, ar: u64, br: u64) -> u64 {
        // mrmul(ar, br) == (ar * br) / r (mod n)
        // R == 2**64
        // gcd(N, R) == 1
        // N * ni mod R == 1
        // 0 <= ar < N < R
        // 0 <= br < N < R
        // T := ar * br
        // t := floor(T / R) - floor(((T * ni mod R) * N) / R)
        // if t < 0 then return t + N else return t
        debug_assert!(ar < self.n());
        debug_assert!(br < self.n());
        let t: u128 = (ar as u128) * (br as u128);
        let (t, f) = ((t >> 64) as u64).overflowing_sub((((((t as u64).wrapping_mul(self.ni())) as u128) * (self.n() as u128)) >> 64) as u64);
        if f { t.wrapping_add(self.n()) } else { t }
    }
    #[inline]
    fn mr(&self, ar: u64) -> u64 {
        // mr(ar) == ar / r (mod n)
        // R == 2**64
        // gcd(N, R) == 1
        // N * ni mod R == 1
        // 0 <= ar < N < R
        // t := floor(ar / R) - floor(((ar * ni mod R) * N) / R)
        // if t < 0 then return t + N else return t
        debug_assert!(ar < self.n());
        let (t, f) = (((((ar.wrapping_mul(self.ni())) as u128) * (self.n() as u128)) >> 64) as u64).overflowing_neg();
        if f { t.wrapping_add(self.n()) } else { t }
    }
    #[inline]
    fn ar(&self, a: u64) -> u64 {
        // ar(a) == a * r (mod n)
        debug_assert!(a < self.n());
        self.mrmul(a, self.r2())
    }
    #[inline]
    fn pow(&self, mut ar: u64, mut b: u64) -> u64 {
        // powir(ar, b) == ((ar / r) ** b) * r (mod n)
        debug_assert!(ar < self.n());
        let mut t = if (b & 1) == 0 { self.r } else { ar };
        loop {
            b >>= 1;
            if b == 0 { return t; }
            ar = self.mrmul(ar, ar);
            if (b & 1) != 0 { t = self.mrmul(t, ar); }
        }
    }
}

#[inline]
fn prime_test_miller_sub(mont: &Mont<u64>, mut base: u64) -> bool {
    // Mirrer-Rabin primality test
    if base >= mont.n() { base %= mont.n(); if base == 0 { return true; } }
    let mut tr = mont.pow(mont.ar(base), mont.d());
    if tr == mont.r() || tr == mont.rn() { return true; }
    for _ in 1..(mont.k()) { tr = mont.mrmul(tr, tr); if tr == mont.rn() { return true; } }
    false
}
pub fn miller_rabin(n: u64) -> bool {
    if n == 2 { return true; }
    if n < 2 || (n & 1) == 0 { return false; }
    let mont = Mont::<u64>::new(n);
    [2,325,9375,28178,450775,9780504,1795265022].iter().all(|&base| prime_test_miller_sub(&mont, base))
}

fn main() {
    use std::io::{BufRead,Write};
    let start_time = std::time::Instant::now();
    let out = std::io::stdout();
    let mut out = std::io::BufWriter::new(out.lock());
    macro_rules! puts {($($format:tt)*) => (let _ = write!(out,$($format)*););}
    let input = std::io::stdin();
    let mut lines = std::io::BufReader::new(input.lock()).lines();
    let n: usize = lines.next().unwrap().unwrap().parse().unwrap();
    for _ in 0..n {
        let x: u64 = lines.next().unwrap().unwrap().parse().unwrap();
        puts!("{} {}\n", x, if miller_rabin(x) { "1" } else { "0" });
    }
    eprint!("{}us\n", start_time.elapsed().as_micros());
}