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

• 問題: No.3030 ミラー・ラビン素数判定法のテスト
• ユーザー: 👑 Mizar
• 提出日時: 2022-08-19 21:07:07
• 言語: Rust (1.77.0 + proconio)
• 結果: AC
• 実行時間: 28 ms / 9,973 ms
• コード長: 7,676 bytes
• コンパイル時間: 13,010 ms
• コンパイル使用メモリ: 378,792 KB
• 実行使用メモリ: 5,248 KB
• 最終ジャッジ日時: 2024-11-16 23:53:27
• 合計ジャッジ時間: 14,096 ms
• ジャッジサーバーID （参考情報）: judge5 / judge3

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 5,248 KB
testcase_01	AC	1 ms; 5,248 KB
testcase_02	AC	1 ms; 5,248 KB
testcase_03	AC	1 ms; 5,248 KB
testcase_04	AC	18 ms; 5,248 KB
testcase_05	AC	18 ms; 5,248 KB
testcase_06	AC	9 ms; 5,248 KB
testcase_07	AC	9 ms; 5,248 KB
testcase_08	AC	9 ms; 5,248 KB
testcase_09	AC	28 ms; 5,248 KB

ソースコード

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

use std::io::Write;

// https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
macro_rules! input {
    ($($r:tt)*) => {
        let stdin = std::io::stdin();
        let mut bytes = std::io::Read::bytes(std::io::BufReader::new(stdin.lock()));
        let mut next = move || -> String{
            bytes.by_ref().map(|r|r.unwrap() as char)
                .skip_while(|c|c.is_whitespace())
                .take_while(|c|!c.is_whitespace())
                .collect()
        };
        input_inner!{next, $($r)*}
    };
}

macro_rules! input_inner {
    ($next:expr) => {};
    ($next:expr,) => {};
    ($next:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($next, $t);
        input_inner!{$next $($r)*}
    };
}

macro_rules! read_value {
    ($next:expr, ( $($t:tt),* )) => { ($(read_value!($next, $t)),*) };
    ($next:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($next, $t)).collect::<Vec<_>>()
    };
    ($next:expr, chars) => {
        read_value!($next, String).chars().collect::<Vec<char>>()
    };
    ($next:expr, usize1) => (read_value!($next, usize) - 1);
    ($next:expr, [ $t:tt ]) => {{
        let len = read_value!($next, usize);
        read_value!($next, [$t; len])
    }};
    ($next:expr, $t:ty) => ($next().parse::<$t>().expect("Parse error"));
}

trait Change { fn chmax(&mut self, x: Self); fn chmin(&mut self, x: Self); }
impl<T: std::cmp::PartialOrd> Change for T {
    fn chmax(&mut self, x: T) { if *self < x { *self = x; } }
    fn chmin(&mut self, x: T) { if *self > x { *self = x; } }
}

fn main() {
    let out = std::io::stdout();
    let mut out = std::io::BufWriter::new(out.lock());
    macro_rules! puts {($($format:tt)*) => (let _ = write!(out,$($format)*););}
    input! {
        n: usize,
        x: [u64; n],
    }
    for i in 0..n {
        puts!("{} {}\n", x[i], if prime_test(x[i]) { 1 } else { 0 });
    }
}

fn prime_test(m: u64) -> bool {
    if m == 2 { return true; }
    if m == 1 || (m & 1) == 0 { return false; }
    let u64mont = U64Mont::new(m);
    match m {
        // Deterministic variants of the Miller-Rabin primality test
        // http://miller-rabin.appspot.com/
        0..=341531 =>
            u64mont.prime_test_once(9345883071009581737),
        0..=1050535501 =>
            u64mont.prime_test_once(336781006125) &&
            u64mont.prime_test_once(9639812373923155),
        0..=350269456337 =>
            u64mont.prime_test_once(4230279247111683200) &&
            u64mont.prime_test_once(14694767155120705706) &&
            u64mont.prime_test_once(16641139526367750375),
        0..=55245642489451 =>
            u64mont.prime_test_once(2) &&
            u64mont.prime_test_once(141889084524735) &&
            u64mont.prime_test_once(1199124725622454117) &&
            u64mont.prime_test_once(11096072698276303650),
        0..=7999252175582851 =>
            u64mont.prime_test_once(2) &&
            u64mont.prime_test_once(4130806001517) &&
            u64mont.prime_test_once(149795463772692060) &&
            u64mont.prime_test_once(186635894390467037) &&
            u64mont.prime_test_once(3967304179347715805),
        0..=585226005592931977 =>
            u64mont.prime_test_once(2) &&
            u64mont.prime_test_once(123635709730000) &&
            u64mont.prime_test_once(9233062284813009) &&
            u64mont.prime_test_once(43835965440333360) &&
            u64mont.prime_test_once(761179012939631437) &&
            u64mont.prime_test_once(1263739024124850375),
        _ =>
            u64mont.prime_test_once(2) &&
            u64mont.prime_test_once(325) &&
            u64mont.prime_test_once(9375) &&
            u64mont.prime_test_once(28178) &&
            u64mont.prime_test_once(450775) &&
            u64mont.prime_test_once(9780504) &&
            u64mont.prime_test_once(1795265022),
    }
}

// モンゴメリ剰余乗算 (Montgomery modular multiplication)
pub trait UMontTrait<T> {
    fn new(m: T) -> Self;
    fn mrmul(&self, ar: T, br: T) -> T; // == (ar * br) / r (mod m)
    fn mr(&self, ar: T) -> T; // == ar / r (mod m)
    fn ar(&self, a: T) -> T; // == a * r (mod m)
    fn powir(&self, ar: T, b: T) -> T; // == ((ar / r) ** b) * r (mod m)
    fn prime_test_once(&self, base: T) -> bool; // Miller-Rabin primality test
}
pub struct U64Mont {
    m: u64, // m is odd, and m > 2
    mi: u64, // m * mi == 1 (mod 2**64)
    r: u64, // == 2**64 (mod m)
    r2: u64, // == 2**128 (mod m)
}
impl UMontTrait<u64> for U64Mont {
    fn new(m: u64) -> Self {
        debug_assert_eq!(m & 1, 1);
        // // m is odd number, m = 2*k+1, m >= 1, m < 2**64, k is non-negative integer, k >= 0, k < 2**63
        // mi0 := m; // = 2*k+1 = (1+(2**2)*((k*(k+1))**1))/(2*k+1)
        let mut mi = m;
        // mi1 := mi0 * (2 - (m * mi0)); // = (1-(2**4)*((k*(k+1))**2))/(2*k+1)
        // mi2 := mi1 * (2 - (m * mi1)); // = (1-(2**8)*((k*(k+1))**4))/(2*k+1)
        // mi3 := mi2 * (2 - (m * mi2)); // = (1-(2**16)*((k*(k+1))**8))/(2*k+1)
        // mi4 := mi3 * (2 - (m * mi3)); // = (1-(2**32)*((k*(k+1))**16))/(2*k+1)
        // mi5 := mi4 * (2 - (m * mi4)); // = (1-(2**64)*((k*(k+1))**32))/(2*k+1)
        // // (m * mi5) mod 2**64 = ((2*k+1) * mi5) mod 2**64 = 1 mod 2**64
        for _ in 0..5 {
            mi = mi.wrapping_mul(2u64.wrapping_sub(m.wrapping_mul(mi)));
        }
        debug_assert_eq!(m.wrapping_mul(mi), 1); // m * mi == 1 (mod 2**64)
        let r: u64 = 0xffff_ffff_ffff_ffff % m + 1; // == 2**64 (mod m)
        let r2: u64 = ((0xffff_ffff_ffff_ffff_ffff_ffff_ffff_ffff % (m as u128)) as u64) + 1; // == 2**128 (mod m)
        debug_assert_eq!(Self { m, mi, r, r2 }.mr(r), 1); // r / r == 1 (mod m)
        debug_assert_eq!(Self { m, mi, r, r2 }.mrmul(1, r2), r); // r2 / r == r (mod m)
        Self { m, mi, r, r2 }
    }
    fn mrmul(&self, ar: u64, br: u64) -> u64 {
        // == (ar * br) / r (mod m)
        debug_assert!(ar < self.m);
        debug_assert!(br < self.m);
        let (m, mi) = (self.m, self.mi);
        let t: u128 = (ar as u128) * (br as u128);
        let t: (u64, bool) = ((t >> 64) as u64).overflowing_sub((((((t as u64).wrapping_mul(mi)) as u128) * (m as u128)) >> 64) as u64);
        if t.1 { t.0.wrapping_add(m) } else { t.0 }
    }
    fn mr(&self, ar: u64) -> u64 {
        // == ar / r (mod m)
        debug_assert!(ar < self.m);
        let (m, mi) = (self.m, self.mi);
        let t: (u64, bool) = (((((ar.wrapping_mul(mi)) as u128) * (m as u128)) >> 64) as u64).overflowing_neg();
        if t.1 { t.0.wrapping_add(m) } else { t.0 }
    }
    fn ar(&self, a: u64) -> u64 {
        // == a * r (mod m)
        debug_assert!(a < self.m);
        self.mrmul(a, self.r2)
    }
    fn powir(&self, mut ar: u64, mut b: u64) -> u64 {
        // == ((ar / r) ** b) * r (mod m)
        debug_assert!(ar < self.m);
        let mut t: u64 = if (b & 1) == 0 { self.r } else { ar };
        b >>= 1;
        while b > 0 {
            ar = self.mrmul(ar, ar);
            t = self.mrmul(t, if (b & 1) == 0 { self.r } else { ar });
            b >>= 1;
        }
        t
    }
    fn prime_test_once(&self, base: u64) -> bool {
        // Miller-Rabin primality test
        debug_assert!(base > 1);
        let (m, r) = (self.m, self.r);
        let mut d = m - 1;
        let k = d.trailing_zeros();
        d >>= k;
        let b = base % m;
        if b == 0 { return true; }
        let mut br = self.powir(self.ar(b), d);
        if br == r { return true; }
        let negr = m - r;
        for _ in 0..k {
            if br == negr { return true; }
            br = self.mrmul(br, br);
        }
        false
    }
}