@@ -29,24 +29,21 @@ pub fn canon_egcd(a: i64, b: i64, c: i64) -> Option<(i64, i64, i64)> {
2929 }
3030}
3131
32+ // TODO: deduplicate modular arithmetic code with num::Field
3233fn pos_mod ( n : i64 , m : i64 ) -> i64 {
3334 if n < 0 {
3435 n + m
3536 } else {
3637 n
3738 }
3839}
39-
4040fn mod_mul ( a : i64 , b : i64 , m : i64 ) -> i64 {
4141 pos_mod ( ( a as i128 * b as i128 % m as i128 ) as i64 , m)
4242}
43-
44- /// Assuming m >= 1 and exp >= 0, finds base ^ exp % m in logarithmic time
45- fn mod_exp ( mut base : i64 , mut exp : i64 , m : i64 ) -> i64 {
43+ fn mod_exp ( mut base : i64 , mut exp : u64 , m : i64 ) -> i64 {
4644 assert ! ( m >= 1 ) ;
47- assert ! ( exp >= 0 ) ;
4845 let mut ans = 1 % m;
49- base = base % m;
46+ base %= m;
5047 while exp > 0 {
5148 if exp % 2 == 1 {
5249 ans = mod_mul ( ans, base, m) ;
@@ -57,12 +54,12 @@ fn mod_exp(mut base: i64, mut exp: i64, m: i64) -> i64 {
5754 pos_mod ( ans, m)
5855}
5956
60- fn is_strong_probable_prime ( n : i64 , d : i64 , r : i64 , a : i64 ) -> bool {
61- let mut x = mod_exp ( a, d , n) ;
57+ fn is_strong_probable_prime ( n : i64 , exp : u64 , r : i64 , a : i64 ) -> bool {
58+ let mut x = mod_exp ( a, exp , n) ;
6259 if x == 1 || x == n - 1 {
6360 return true ;
6461 }
65- for _ in 0 .. ( r - 1 ) {
62+ for _ in 1 ..r {
6663 x = mod_mul ( x, x, n) ;
6764 if x == n - 1 {
6865 return true ;
@@ -71,21 +68,22 @@ fn is_strong_probable_prime(n: i64, d: i64, r: i64, a: i64) -> bool {
7168 false
7269}
7370
74- const BASES : [ i64 ; 12 ] = [ 2 , 3 , 5 , 7 , 11 , 13 , 17 , 19 , 23 , 29 , 31 , 37 ] ;
75- /// Assuming x >= 0, returns true if x is prime
71+ /// Assuming x >= 0, returns whether x is prime
7672pub fn is_prime ( n : i64 ) -> bool {
73+ const BASES : [ i64 ; 12 ] = [ 2 , 3 , 5 , 7 , 11 , 13 , 17 , 19 , 23 , 29 , 31 , 37 ] ;
7774 assert ! ( n >= 0 ) ;
7875 match n {
79- 0 | 1 => return false ,
80- 2 | 3 => return true ,
81- _ if n % 2 == 0 => return false ,
82- _ => { }
83- } ;
84- let r = ( n - 1 ) . trailing_zeros ( ) as i64 ;
85- let d = ( n - 1 ) >> r;
86- BASES
87- . iter ( )
88- . all ( |& base| base > n - 2 || is_strong_probable_prime ( n, d, r, base) )
76+ 0 | 1 => false ,
77+ 2 | 3 => true ,
78+ _ if n % 2 == 0 => false ,
79+ _ => {
80+ let r = ( n - 1 ) . trailing_zeros ( ) as i64 ;
81+ let exp = ( n - 1 ) as u64 >> r;
82+ BASES
83+ . iter ( )
84+ . all ( |& base| base > n - 2 || is_strong_probable_prime ( n, exp, r, base) )
85+ }
86+ }
8987}
9088
9189fn pollard_rho ( n : i64 ) -> i64 {
@@ -96,27 +94,26 @@ fn pollard_rho(n: i64) -> i64 {
9694 loop {
9795 x = f ( x) ;
9896 y = f ( f ( y) ) ;
99- let p = num:: fast_gcd ( x - y, n) ;
100- if p > 1 {
101- if p == n {
102- break ;
103- } else {
104- return p;
105- }
97+ let div = num:: fast_gcd ( x - y, n) ;
98+ if div == n {
99+ break ;
100+ } else if div > 1 {
101+ return div;
106102 }
107103 }
108104 }
109105 panic ! ( "No divisor found!" ) ;
110106}
111107
112108/// Assuming x >= 1, finds the prime factorization of n
109+ /// TODO: pollard_rho needs randomization to ensure correctness in contest settings!
113110pub fn factorize ( n : i64 ) -> Vec < i64 > {
114111 assert ! ( n >= 1 ) ;
115- let r = n. trailing_zeros ( ) ;
116- let mut factors = vec ! [ 2 ; r as usize ] ;
112+ let r = n. trailing_zeros ( ) as usize ;
113+ let mut factors = vec ! [ 2 ; r] ;
117114 let mut stack = match n >> r {
118- 1 => Vec :: new ( ) ,
119- x => vec ! [ x]
115+ 1 => vec ! [ ] ,
116+ x => vec ! [ x] ,
120117 } ;
121118 while let Some ( top) = stack. pop ( ) {
122119 if is_prime ( top) {
@@ -171,7 +168,7 @@ mod test {
171168
172169 #[ test]
173170 fn test_pollard ( ) {
174- assert_eq ! ( factorize( 1 ) , Vec :: new ( ) ) ;
171+ assert_eq ! ( factorize( 1 ) , vec! [ ] ) ;
175172 assert_eq ! ( factorize( 2 ) , vec![ 2 ] ) ;
176173 assert_eq ! ( factorize( 4 ) , vec![ 2 , 2 ] ) ;
177174 assert_eq ! ( factorize( 12 ) , vec![ 2 , 2 , 3 ] ) ;
0 commit comments