competitive/math/
primitive_root.rs

1use super::{BarrettReduction, 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    (2..)
11        .find(|&g| check_primitive_root(g, phi, &br, &pf))
12        .unwrap()
13}
14
15pub fn check_primitive_root(
16    g: u64,
17    phi: u64,
18    br: &BarrettReduction<u128>,
19    pf: &[(u64, u32)],
20) -> bool {
21    pf.iter().all(|&(q, _)| {
22        let mut g = g as u128;
23        let mut k = phi / q;
24        let mut r: u128 = 1;
25        while k > 0 {
26            if k & 1 == 1 {
27                r = br.rem(r * g);
28            }
29            g = br.rem(g * g);
30            k >>= 1;
31        }
32        r > 1
33    })
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39    use crate::math::PrimeList;
40
41    #[test]
42    fn test_primitive_root() {
43        assert_eq!(3, primitive_root(998244353));
44        let pl = PrimeList::new(1000);
45        for &p in pl.primes() {
46            let g = primitive_root(p);
47            let mut x = g;
48            for _ in 1..p - 1 {
49                assert_ne!(x, 1);
50                x = x * g % p;
51            }
52            assert_eq!(x, 1);
53        }
54    }
55
56    #[test]
57    fn test_primitive_root_prime_power() {
58        let pl = PrimeList::new(100);
59        for &p in pl.primes().iter().skip(1) {
60            for e in 1.. {
61                let n = p.pow(e);
62                let phi = n - n / p;
63                let g = {
64                    let pf = prime_factors(phi);
65                    let br = BarrettReduction::<u128>::new(n as _);
66                    (2..)
67                        .find(|&g| check_primitive_root(g, phi, &br, &pf))
68                        .unwrap()
69                };
70                let mut x = g;
71                for _ in 1..phi {
72                    assert_ne!(x, 1);
73                    x = x * g % n;
74                }
75                assert_eq!(x, 1);
76                if n >= 10_000 {
77                    break;
78                }
79            }
80        }
81    }
82
83    #[test]
84    fn test_primitive_root_power_of_two() {
85        for e in 3.. {
86            let n = 2u64.pow(e);
87            let phi = n / 4;
88            let g = 5;
89            let mut x = g;
90            for _ in 1..phi {
91                assert_ne!(x, 1);
92                x = x * g % n;
93            }
94            assert_eq!(x, 1);
95            if n >= 10_000 {
96                break;
97            }
98        }
99    }
100}
competitive/math/primitive_root.rs

competitive/math/
primitive_root.rs
