@@ -29,24 +29,21 @@ pub fn canon_egcd(a: i64, b: i64, c: i64) -> Option<(i64, i64, i64)> {
29
29
}
30
30
}
31
31
32
+ // TODO: deduplicate modular arithmetic code with num::Field
32
33
fn pos_mod ( n : i64 , m : i64 ) -> i64 {
33
34
if n < 0 {
34
35
n + m
35
36
} else {
36
37
n
37
38
}
38
39
}
39
-
40
40
fn mod_mul ( a : i64 , b : i64 , m : i64 ) -> i64 {
41
41
pos_mod ( ( a as i128 * b as i128 % m as i128 ) as i64 , m)
42
42
}
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 {
46
44
assert ! ( m >= 1 ) ;
47
- assert ! ( exp >= 0 ) ;
48
45
let mut ans = 1 % m;
49
- base = base % m;
46
+ base %= m;
50
47
while exp > 0 {
51
48
if exp % 2 == 1 {
52
49
ans = mod_mul ( ans, base, m) ;
@@ -57,12 +54,12 @@ fn mod_exp(mut base: i64, mut exp: i64, m: i64) -> i64 {
57
54
pos_mod ( ans, m)
58
55
}
59
56
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) ;
62
59
if x == 1 || x == n - 1 {
63
60
return true ;
64
61
}
65
- for _ in 0 .. ( r - 1 ) {
62
+ for _ in 1 ..r {
66
63
x = mod_mul ( x, x, n) ;
67
64
if x == n - 1 {
68
65
return true ;
@@ -71,21 +68,22 @@ fn is_strong_probable_prime(n: i64, d: i64, r: i64, a: i64) -> bool {
71
68
false
72
69
}
73
70
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
76
72
pub fn is_prime ( n : i64 ) -> bool {
73
+ const BASES : [ i64 ; 12 ] = [ 2 , 3 , 5 , 7 , 11 , 13 , 17 , 19 , 23 , 29 , 31 , 37 ] ;
77
74
assert ! ( n >= 0 ) ;
78
75
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
+ }
89
87
}
90
88
91
89
fn pollard_rho ( n : i64 ) -> i64 {
@@ -96,27 +94,26 @@ fn pollard_rho(n: i64) -> i64 {
96
94
loop {
97
95
x = f ( x) ;
98
96
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;
106
102
}
107
103
}
108
104
}
109
105
panic ! ( "No divisor found!" ) ;
110
106
}
111
107
112
108
/// Assuming x >= 1, finds the prime factorization of n
109
+ /// TODO: pollard_rho needs randomization to ensure correctness in contest settings!
113
110
pub fn factorize ( n : i64 ) -> Vec < i64 > {
114
111
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] ;
117
114
let mut stack = match n >> r {
118
- 1 => Vec :: new ( ) ,
119
- x => vec ! [ x]
115
+ 1 => vec ! [ ] ,
116
+ x => vec ! [ x] ,
120
117
} ;
121
118
while let Some ( top) = stack. pop ( ) {
122
119
if is_prime ( top) {
@@ -171,7 +168,7 @@ mod test {
171
168
172
169
#[ test]
173
170
fn test_pollard ( ) {
174
- assert_eq ! ( factorize( 1 ) , Vec :: new ( ) ) ;
171
+ assert_eq ! ( factorize( 1 ) , vec! [ ] ) ;
175
172
assert_eq ! ( factorize( 2 ) , vec![ 2 ] ) ;
176
173
assert_eq ! ( factorize( 4 ) , vec![ 2 , 2 ] ) ;
177
174
assert_eq ! ( factorize( 12 ) , vec![ 2 , 2 , 3 ] ) ;
0 commit comments