competitive/algorithm/
baby_step_giant_step.rs

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