competitive/algorithm/
baby_step_giant_step.rs1use super::Monoid;
2use std::{collections::HashSet, hash::Hash};
3
4pub 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}