competitive/math/
primitive_root.rs

1use super::{BarrettReduction, Xorshift, prime_factors};
2
3pub fn primitive_root(p: u64) -> u64 {
4    if p == 2 {
5        return 1;
6    }
7    let phi = p - 1;
8    let pf = prime_factors(phi);
9    let br = BarrettReduction::<u128>::new(p as _);
10    for g in 2..=3.min(p - 1) {
11        if check_primitive_root(g, phi, &br, &pf) {
12            return g;
13        }
14    }
15    let mut rng = Xorshift::default();
16    loop {
17        let g = ((rng.rand64() as u128 * (p - 2) as u128) >> 64) as u64 + 2;
18        if check_primitive_root(g, phi, &br, &pf) {
19            return g;
20        }
21    }
22}
23
24pub fn check_primitive_root(
25    g: u64,
26    phi: u64,
27    br: &BarrettReduction<u128>,
28    pf: &[(u64, u32)],
29) -> bool {
30    pf.iter().all(|&(q, _)| {
31        let mut g = g as u128;
32        let mut k = phi / q;
33        let mut r: u128 = 1;
34        while k > 0 {
35            if k & 1 == 1 {
36                r = br.rem(r * g);
37            }
38            g = br.rem(g * g);
39            k >>= 1;
40        }
41        r > 1
42    })
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48    use crate::math::PrimeList;
49
50    #[test]
51    fn test_primitive_root() {
52        let primes: Vec<_> = PrimeList::new(1000).primes().map(u64::from).collect();
53        for p in primes {
54            let g = primitive_root(p);
55            let mut powers = vec![false; p as usize];
56            let mut x = 1;
57            for _ in 0..p - 1 {
58                assert!(!powers[x as usize]);
59                powers[x as usize] = true;
60                x = x * g % p;
61            }
62            assert_eq!(x, 1);
63            assert!(powers[1..].iter().all(|&seen| seen));
64        }
65    }
66
67    #[test]
68    fn test_check_primitive_root() {
69        let mut rng = Xorshift::default();
70        let primes: Vec<_> = PrimeList::new(30).primes().skip(1).map(u64::from).collect();
71        for _ in 0..1000 {
72            let p = primes[rng.random(0..primes.len())];
73            let exponent = rng.random(1..=3);
74            let n = p.pow(exponent);
75            let phi = n - n / p;
76            let factors = prime_factors(phi);
77            let br = BarrettReduction::<u128>::new(n as _);
78            let g = rng.random(1..n);
79            if g % p == 0 {
80                continue;
81            }
82            let mut x = 1;
83            let order = (1..=phi)
84                .find(|_| {
85                    x = x * g % n;
86                    x == 1
87                })
88                .unwrap();
89            assert_eq!(check_primitive_root(g, phi, &br, &factors), order == phi);
90        }
91    }
92}
competitive/math/primitive_root.rs

competitive/math/
primitive_root.rs
