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

提出ソース
結果

問題	No.3030 ミラー・ラビン素数判定法のテスト
ユーザー	👑 Mizar
提出日時	2022-08-17 15:09:11
言語	Rust (1.77.0)
結果	AC
実行時間	35 ms / 9,973 ms
コード長	12,653 bytes
コンパイル時間	11,585 ms
コンパイル使用メモリ	378,456 KB
実行使用メモリ	6,944 KB
最終ジャッジ日時	2024-04-28 09:52:47
合計ジャッジ時間	12,431 ms
ジャッジサーバーID （参考情報）	judge4 / judge1
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	0 ms 6,812 KB
testcase_01	AC	1 ms 6,940 KB
testcase_02	AC	1 ms 6,940 KB
testcase_03	AC	1 ms 6,944 KB
testcase_04	AC	21 ms 6,940 KB
testcase_05	AC	21 ms 6,940 KB
testcase_06	AC	12 ms 6,940 KB
testcase_07	AC	12 ms 6,940 KB
testcase_08	AC	12 ms 6,940 KB
testcase_09	AC	35 ms 6,940 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
// -*- coding:utf-8-unix -*-

#[allow(unused_imports)]
use std::io::{BufReader, BufWriter, Write, stdin, stdout, stderr};
#[allow(unused_imports)]
use std::collections::*;
#[allow(unused_imports)]
use std::cmp::*;
#[allow(unused_imports)]
use std::collections::*;

// 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: 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 = BufWriter::with_capacity(0x4_0000, 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; }
    match m {
        0..=49141 => {
            let umont = U32Mont::new(m as u32);
            umont.prime_test_once(921211727)
        },
        0..=316349281 => {
            let umont = U32Mont::new(m as u32);
            umont.prime_test_once(11000544) &&
            umont.prime_test_once(31481107)
        },
        0..=0x7fffffff => {
            let umont = U32Mont::new(m as u32);
            umont.prime_test_once(2) &&
            umont.prime_test_once(7) &&
            umont.prime_test_once(61)
        },
        0..=350269456337 => {
            let umont = U64Mont::new(m);
            umont.prime_test_once(4230279247111683200) &&
            umont.prime_test_once(14694767155120705706) &&
            umont.prime_test_once(16641139526367750375)
        },
        0..=55245642489451 => {
            let umont = U64Mont::new(m);
            umont.prime_test_once(2) &&
            umont.prime_test_once(141889084524735) &&
            umont.prime_test_once(1199124725622454117) &&
            umont.prime_test_once(11096072698276303650)
        },
        0..=7999252175582851 => {
            let umont = U64Mont::new(m);
            umont.prime_test_once(2) &&
            umont.prime_test_once(4130806001517) &&
            umont.prime_test_once(149795463772692060) &&
            umont.prime_test_once(186635894390467037) &&
            umont.prime_test_once(3967304179347715805)
        },
        0..=585226005592931977 => {
            let umont = U64Mont::new(m);
            umont.prime_test_once(2) &&
            umont.prime_test_once(123635709730000) &&
            umont.prime_test_once(9233062284813009) &&
            umont.prime_test_once(43835965440333360) &&
            umont.prime_test_once(761179012939631437) &&
            umont.prime_test_once(1263739024124850375)
        },
        _ => {
            let umont = U64Mont::new(m);
            umont.prime_test_once(2) &&
            umont.prime_test_once(325) &&
            umont.prime_test_once(9375) &&
            umont.prime_test_once(28178) &&
            umont.prime_test_once(450775) &&
            umont.prime_test_once(9780504) &&
            umont.prime_test_once(1795265022)
        },
    }
}

pub trait UMontTrait<T> {
    fn new(m: T) -> Self;
    fn mrmul(&self, ar: T, br: T) -> T; // == (ar * br) / r (mod m)
    fn mrmuli(&self, ar: T, br: T) -> T; // == (ar * br) / r (mod m) (m < T::max_value()/2)
    fn mr(&self, ar: T) -> T; // == ar / r (mod m)
    fn ar(&self, a: T) -> T; // == a * r (mod m)
    fn ari(&self, a: T) -> T; // == a * r (mod m) (m < T::max_value()/2)
    fn powir(&self, ar: T, b: T) -> T; // == ((ar / r) ** b) * r (mod m)
    fn powiri(&self, ar: T, b: T) -> T; // == ((ar / r) ** b) * r (mod m) (m < T::max_value()/2)
    fn prime_test_once(&self, base: T) -> bool;
}
pub struct U32Mont {
    m: u32, // is odd
    md: u32, // m * md + 1 == 0 (mod 2**32)
    r: u32, // == 2**32 (mod m)
    r2: u32, // == 2**64 (mod m)
}
impl UMontTrait<u32> for U32Mont {
    fn new(m: u32) -> Self {
        debug_assert!(m > 2);
        debug_assert!((m & 1) == 1);
        let mut t: u32 = 0;
        let mut md: u32 = 0;
        for i in 0..32 {
            if (t & 1) == 0 {
                t = t.wrapping_add(m);
                md |= 1 << i;
            }
            t >>= 1;
        }
        let r = m.wrapping_neg() % m;
        let r2 = (((r as u64) * (r as u64)) % (m as u64)) as u32;
        debug_assert_eq!(Self { m, md, r, r2 }.mr(r), 1);
        debug_assert_eq!(Self { m, md, r, r2 }.mrmul(1, r2), r);
        Self { m, md, r, r2 }
    }
    fn mrmul(&self, ar: u32, br: u32) -> u32 {
        debug_assert!(ar < self.m);
        debug_assert!(br < self.m);
        let (m, md) = (self.m, self.md);
        let t = (ar as u64) * (br as u64);
        // 2**31 < m にも対応
        let t = (t >> 32) + (((t & 0xffff_ffff) + ((((t as u32).wrapping_mul(md)) as u64) * (m as u64))) >> 32);
        let r = t.overflowing_sub(m as u64);
        if r.1 { t as u32 } else { r.0 as u32 }
    }
    fn mrmuli(&self, ar: u32, br: u32) -> u32 {
        debug_assert!(ar < self.m);
        debug_assert!(br < self.m);
        let (m, md) = (self.m, self.md);
        let t = (ar as u64) * (br as u64);
        // m < 2**31 のみ対応
        let t = ((t + ((((t as u32).wrapping_mul(md)) as u64) * (m as u64))) >> 32) as u32;
        if t >= m { t - m } else { t }
    }
    fn mr(&self, ar: u32) -> u32 {
        debug_assert!(ar < self.m);
        let (m, md) = (self.m, self.md);
        let t = (((ar as u64) + (((ar.wrapping_mul(md)) as u64) * (m as u64))) >> 32) as u32;
        if t >= m { t - m } else { t }
    }
    fn ar(&self, a: u32) -> u32 {
        self.mrmul(a, self.r2)
    }
    fn ari(&self, a: u32) -> u32 {
        self.mrmuli(a, self.r2)
    }
    fn powir(&self, mut ar: u32, mut b: u32) -> u32 {
        debug_assert!(ar < self.m);
        let mut res = if (b & 1) == 1 { ar } else { self.r };
        b >>= 1;
        while b > 0 {
            ar = self.mrmul(ar, ar);
            if (b & 1) == 1 {
                res = self.mrmul(res, ar);
            }
            b >>= 1;
        }
        res
    }
    fn powiri(&self, mut ar: u32, mut b: u32) -> u32 {
        debug_assert!(ar < self.m);
        let mut res = if (b & 1) == 1 { ar } else { self.r };
        b >>= 1;
        while b > 0 {
            ar = self.mrmuli(ar, ar);
            if (b & 1) == 1 {
                res = self.mrmuli(res, ar);
            }
            b >>= 1;
        }
        res
    }
    fn prime_test_once(&self, base: u32) -> bool {
        let mut d = self.m - 1;
        let k = d.trailing_zeros();
        d >>= k;
        match self.m {
            0..=0x7fff_ffff => {
                let mut br = self.ari(base % self.m);
                if br == 0 { return true; }
                br = self.powiri(br, d);
                if br == self.r { return true; }
                let negr = self.m - self.r;
                for _ in 0..k {
                    if br == negr { return true; }
                    br = self.mrmuli(br, br);
                }
                false
            },
            _ => {
                let mut br = self.ar(base % self.m);
                if br == 0 { return true; }
                br = self.powir(br, d);
                if br == self.r { return true; }
                let negr = self.m - self.r;
                for _ in 0..k {
                    if br == negr { return true; }
                    br = self.mrmul(br, br);
                }
                false
            },
        }
    }
}
pub struct U64Mont {
    m: u64, // is odd
    md: u64, // m * md + 1 == 0 (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!(m > 2);
        debug_assert!((m & 1) == 1);
        let mut t: u64 = 0;
        let mut md: u64 = 0;
        for i in 0..64 {
            if (t & 1) == 0 {
                t = t.wrapping_add(m);
                md |= 1 << i;
            }
            t >>= 1;
        }
        let r = m.wrapping_neg() % m;
        let r2 = (((r as u128) * (r as u128)) % (m as u128)) as u64;
        debug_assert_eq!(Self { m, md, r, r2 }.mr(r), 1);
        debug_assert_eq!(Self { m, md, r, r2 }.mrmul(1, r2), r);
        Self { m, md, r, r2 }
    }
    fn mrmul(&self, ar: u64, br: u64) -> u64 {
        debug_assert!(ar < self.m);
        debug_assert!(br < self.m);
        let (m, md) = (self.m, self.md);
        let t = (ar as u128) * (br as u128);
        // m > 2**63 にも対応
        let t = (t >> 64) + (((t & 0xffff_ffff_ffff_ffff) + ((((t as u64).wrapping_mul(md)) as u128) * (m as u128))) >> 64);
        let r = t.overflowing_sub(m as u128);
        if r.1 { t as u64 } else { r.0 as u64 }
    }
    fn mrmuli(&self, ar: u64, br: u64) -> u64 {
        debug_assert!(ar < self.m);
        debug_assert!(br < self.m);
        let (m, md) = (self.m, self.md);
        let t = (ar as u128) * (br as u128);
        // m < 2**63 のみ対応
        let t = ((t + ((((t as u64).wrapping_mul(md)) as u128) * (m as u128))) >> 64) as u64;
        if t >= m { t - m } else { t }
    }
    fn mr(&self, ar: u64) -> u64 {
        debug_assert!(ar < self.m);
        let (m, md) = (self.m, self.md);
        let t = (((ar as u128) + (((ar.wrapping_mul(md)) as u128) * (m as u128))) >> 64) as u64;
        if t >= m { t - m } else { t }
    }
    fn ar(&self, a: u64) -> u64 {
        self.mrmul(a, self.r2)
    }
    fn ari(&self, a: u64) -> u64 {
        self.mrmuli(a, self.r2)
    }
    fn powir(&self, mut ar: u64, mut b: u64) -> u64 {
        debug_assert!(ar < self.m);
        let mut res = if (b & 1) == 1 { ar } else { self.r };
        b >>= 1;
        while b > 0 {
            ar = self.mrmul(ar, ar);
            if (b & 1) == 1 {
                res = self.mrmul(res, ar);
            }
            b >>= 1;
        }
        res
    }
    fn powiri(&self, mut ar: u64, mut b: u64) -> u64 {
        debug_assert!(ar < self.m);
        let mut res = if (b & 1) == 1 { ar } else { self.r };
        b >>= 1;
        while b > 0 {
            ar = self.mrmuli(ar, ar);
            if (b & 1) == 1 {
                res = self.mrmuli(res, ar);
            }
            b >>= 1;
        }
        res
    }
    fn prime_test_once(&self, base: u64) -> bool {
        let mut d = self.m - 1;
        let k = d.trailing_zeros();
        d >>= k;
        match self.m {
            0..=0x7fff_ffff_ffff_ffff => {
                let mut b = self.ari(base % self.m);
                if b == 0 { return true; }
                b = self.powiri(b, d);
                if b == self.r { return true; }
                let negr = self.m - self.r;
                for _ in 0..k {
                    if b == negr { return true; }
                    b = self.mrmuli(b, b);
                }
                false
            },
            _ => {
                let mut b = self.ar(base % self.m);
                if b == 0 { return true; }
                b = self.powir(b, d);
                if b == self.r { return true; }
                let negr = self.m - self.r;
                for _ in 0..k {
                    if b == negr { return true; }
                    b = self.mrmul(b, b);
                }
                false
            },
        }
    }
}
yukicoder

結果

テストケース

ソースコード
