competitive/algorithm/
baby_step_giant_step.rs

1use super::Monoid;
2use std::{collections::HashSet, hash::Hash};
3
4/// $\min\{0\le i < n | a x^i=b\}$
5pub fn baby_step_giant_step<M>(a: M::T, x: M::T, b: M::T, n: usize) -> Option<usize>
6where
7    M: Monoid<T: Eq + Hash>,
8{
9    if a == b {
10        return Some(0);
11    }
12    let block_size = 1usize.max((n as f64).sqrt() as _);
13    let mut baby = HashSet::new();
14    let mut t = b.clone();
15    for _ in 0..block_size {
16        t = M::operate(&t, &x);
17        baby.insert(t.clone());
18    }
19    let g = M::pow(x.clone(), block_size);
20    let mut t = a;
21    let mut fail = 0usize;
22    for k in (0..n).step_by(block_size) {
23        let nt = M::operate(&t, &g);
24        if baby.contains(&nt) {
25            for m in k..n.min(k + block_size) {
26                if t == b {
27                    return Some(m);
28                }
29                t = M::operate(&t, &x);
30            }
31            fail += 1;
32            if fail >= 2 {
33                break;
34            }
35        }
36        t = nt;
37    }
38    None
39}
40
41#[cfg(test)]
42mod tests {
43    use super::*;
44    use crate::{
45        algebra::MultiplicativeOperation as MulOp, num::mint_basic::DynMIntU32, tools::Xorshift,
46    };
47
48    #[test]
49    fn test_bsgs_small() {
50        for n in 2..=32 {
51            DynMIntU32::set_mod(n);
52            for a in 0..n {
53                for x in 0..n {
54                    let (a, x) = (DynMIntU32::new(a), DynMIntU32::new(x));
55                    let powers: Vec<_> = (0..n)
56                        .scan(a, |value, _| {
57                            let current = *value;
58                            *value *= x;
59                            Some(current)
60                        })
61                        .collect();
62                    for b in 0..n {
63                        let b = DynMIntU32::new(b);
64                        let expected = powers.iter().position(|&value| value == b);
65                        assert_eq!(
66                            baby_step_giant_step::<MulOp<DynMIntU32>>(a, x, b, n as _),
67                            expected
68                        );
69                    }
70                }
71            }
72        }
73    }
74
75    #[test]
76    fn test_bsgs_medium() {
77        let mut rng = Xorshift::default();
78        for _ in 0..100 {
79            let n = rng.random(2..100_000u32);
80            DynMIntU32::set_mod(n);
81            let a = DynMIntU32::new(rng.random(..n));
82            let x = DynMIntU32::new(rng.random(..n));
83            let b = DynMIntU32::new(rng.random(..n));
84            let mut value = a;
85            let exp = (0..n).position(|_| {
86                let found = value == b;
87                value *= x;
88                found
89            });
90            let ans = baby_step_giant_step::<MulOp<DynMIntU32>>(a, x, b, n as _);
91            assert_eq!(exp, ans);
92        }
93    }
94
95    #[test]
96    fn test_bsgs_large() {
97        let mut rng = Xorshift::default();
98        for _ in 0..100 {
99            let n = rng.random(2..1_000_000_000u32);
100            DynMIntU32::set_mod(n);
101            let a = DynMIntU32::new(rng.random(..n));
102            let x = DynMIntU32::new(rng.random(..n));
103            let exponent = rng.random(0..n);
104            let b = a * x.pow(exponent as _);
105            let ans = baby_step_giant_step::<MulOp<DynMIntU32>>(a, x, b, n as _);
106            let i = ans.unwrap();
107            assert_eq!(a * x.pow(i), b);
108            assert!(i < n as usize);
109        }
110    }
111}
competitive/algorithm/baby_step_giant_step.rs

competitive/algorithm/
baby_step_giant_step.rs
