competitive/algorithm/
baby_step_giant_step.rs

1use super::Monoid;
2use std::{collections::HashSet, hash::Hash};
3
4/// $\min\{0\le i < n | x^i=y\}$
5pub fn baby_step_giant_step<M>(x: M::T, y: M::T, n: usize) -> Option<usize>
6where
7    M: Monoid<T: Eq + Hash>,
8{
9    if M::is_unit(&y) {
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 = y.clone();
15    for _ in 0..block_size {
16        t = M::operate(&x, &t);
17        baby.insert(t.clone());
18    }
19    let g = M::pow(x.clone(), block_size);
20    let mut t = M::unit();
21    let mut fail = 0usize;
22    for k in (0..n).step_by(block_size) {
23        let nt = M::operate(&g, &t);
24        if baby.contains(&nt) {
25            for m in k..n.min(k + block_size) {
26                if t == y {
27                    return Some(m);
28                }
29                t = M::operate(&x, &t);
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..50 {
51            DynMIntU32::set_mod(n);
52            for x in 0..n {
53                for y in 0..n {
54                    let (x, y) = (DynMIntU32::new(x), DynMIntU32::new(y));
55                    let exp = (0..n).position(|i| x.pow(i as _) == y);
56                    let ans = baby_step_giant_step::<MulOp<DynMIntU32>>(x, y, n as _);
57                    assert_eq!(exp, ans);
58                }
59            }
60        }
61    }
62
63    #[test]
64    fn test_bsgs_midium() {
65        let mut rng = Xorshift::default();
66        for _ in 0..10 {
67            let n = rng.random(2..100_000u32);
68            DynMIntU32::set_mod(n);
69            let x = DynMIntU32::new(rng.random(..n));
70            let y = DynMIntU32::new(rng.random(..n));
71            let exp = (0..n).position(|i| x.pow(i as _) == y);
72            let ans = baby_step_giant_step::<MulOp<DynMIntU32>>(x, y, n as _);
73            assert_eq!(exp, ans);
74        }
75    }
76
77    #[test]
78    fn test_bsgs_large() {
79        let mut rng = Xorshift::default();
80        for _ in 0..20 {
81            let n = rng.random(2..1_000_000_000u32);
82            DynMIntU32::set_mod(n);
83            let x = DynMIntU32::new(rng.random(..n));
84            let y = DynMIntU32::new(rng.random(..n));
85            let ans = baby_step_giant_step::<MulOp<DynMIntU32>>(x, y, n as _);
86            if let Some(i) = ans {
87                assert_eq!(x.pow(i), y);
88                assert!(i < n as usize);
89            }
90        }
91    }
92}