Baby step Giant stepを完全に理解してABC270のG問題をRustで解く

はじめに

これはABC270のG問題の解説を読んでRustで実装したものです。

Baby step Giant step について

Wikipediaなどで調べると、Baby step Giant stepは離散対数問題を解くアルゴリズムのようです。
素数$p$と$x, y$が与えられたとき、「$x^n\equiv y\mod p$を満たす$n$が存在するかどうか」を判定し、存在する場合は最小の$n$を求めることができます。

そのアルゴリズムは以下です。

1. 自然数$m$を決める。
2. （Baby step）$S=\lbrace y, y(x^{-1}), y(x^{-1})^2,\ldots,y(x^{-1})^{m-1}\rbrace$を求める。
3. (Giant step)$i=0,\ldots,\lfloor \frac{p}{m}\rfloor$ の順に以下を行う。
   1. $x^{im}$が$S$に含まれているか判定する。
      • $x^{im}$の計算はBaby step終了後に$x^m$を求めておくことで、$x^{im}=x^{(i-1)m}\cdot x^m$とすると効率的です。
   2. 含まれている場合は$x^{im}=y(x^{-1})^j$であるような最小の$j$を求める。$n=im+j$が条件を満たすので、ループを打ち切って終了する。

これは空間計算量$O(m)$、時間計算量$O(m+\frac{p}{m})$です。$m=\lfloor\sqrt{p}\rfloor$あたりで時間計算量が最小になります。

さて、このアルゴリズムを応用すると以下の問題も解くことができます。（こちらの方がABCの解説に近いです）

状態$s, g$および次の状態を返す関数$f(x)$が与えられる。$f^n(x)=y$を満たす$n\in[0, N]$が存在するか判定し、存在するならそのような$n$の最小値を求めよ。
ただし、$f$は以下の制約を満たしている。

• $f^m(x)$を高速に計算できる。
  • この計算に$O(m)$かかるようであれば全探索と効率が変わりません。
• $f^{-1}(x)$を求められる。
  • $f^{-1}(x)$がわからないとBaby step で$S$を構築できません。

実装

/// xからyにするにはfを何回適用すればいいか求める
/// # Arguments
///
/// * `x` - 開始状態
/// * `y` - 終了状態
/// * `max` - `f`を適用する回数の上限
/// * `m`, `f_m` - `x`に`f`を`m`回適用したときの状態
/// * `f_inv` - `f`の逆関数
///
/// # Returns
/// * `Some(n)` - `f^n(x)=y`であるような最小の`n`
/// * `None` - 'max`回以下では`x`を`y`にできない
///
fn baby_step_giant_step<T, FM, FINV>(x: T, y: T, max: usize, m: usize, f_m: FM, f_inv: FINV) -> Option<usize>
where
	T: Copy + Eq + std::hash::Hash,
	FM: Fn(T) -> T,
	FINV: Fn(T) -> T
{
	let mut result = None;
	let mut candidates = std::collections::HashSet::new();
	let mut candidate = y;
	for _j in (0..m).rev() {
		candidates.insert(candidate);
		candidate = f_inv(candidate);
	}

	let mut fmx = x;
	for i in 0..=max/m {
		if candidates.contains(&fmx) {
			let mut v = f_m(fmx);
			for j in (0..m).rev() {
				v = f_inv(v);
				if v == y {
					result = Some(i * m + j);
				}
			}

            break;
		}

        fmx = f_m(fmx);
	}

	result
}

実際にABC270のG問題をACした解答はこちら