Submission #35152746

---

Source Code Expand

use proconio::input;

fn main() {
    input! {
        t: usize,
        case: [(usize, usize, usize, usize, usize); t]
    }

    for (p, a, b, s, g) in case {
        if a == 0 {
            if s == g {
                println!("0");
            } else if b == g {
                println!("1");
            } else {
                println!("-1");
            }
        } else {
            // f(x) = (ax+b) % p
            // fm(x) = (a^mx + a^(m-1)b + ... + b) % p = (a^mx + b * (a^m-1) /  (a-1)) % p
            let m = sqrt(p);
            let a_inv = mod_div(1, a, p);

            let f_inv = |y: usize| {
                (y + p - b) *  a_inv %  p
            };

            let result = if a == 1 {
                let f_m = |x: usize| {
                    (x + m * b) % p
                };
                baby_step_giant_step(s, g, p, m, f_m, f_inv)
            } else {
                let a_pow_m = mod_pow(a, m, p);
                let a_minus_1_inv = mod_div(1, a-1, p);
                let f_m = |x: usize| {
                    (a_pow_m * x + b * ((a_pow_m-1) * a_minus_1_inv % p)) % p
                };
                baby_step_giant_step(s, g, p, m, f_m, f_inv)
            };

            match result {
                Some(result) => {
                    println!("{}", result);
                }
                None => {
                    println!("-1");
                }
            }
        }
    }
}



/// 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 f_im_x = x;
	for i in 0..=max/m {
		if candidates.contains(&f_im_x) {
			let mut v = f_m(f_im_x);
			for j in (0..m).rev() {
				v = f_inv(v);
				if v == y {
					result = Some(i * m + j);
				}
			}

            break;
		}

        f_im_x = f_m(f_im_x);
	}

	result
}


/// 平方根
///
/// # Arguments
///
/// * `n` - 対象の数
///
/// # Returns
///
/// n^0.5（小数点以下切り捨て）
///
pub fn sqrt(n: usize) -> usize {
	if n == 0 {
		0
	} else {
		b_search_find_last(1, n, |i| i <= n / i).unwrap()
	}
}

pub fn mod_pow(n: usize, mut k: usize, p: usize) -> usize {
    let mut val = n;
    let mut result: usize = 1;
    while k > 0 {
        if k % 2 == 1 {
            result = (result * val) % p;
        }
        val = (val * val) % p;
        k /= 2;
    }

    result
}

pub fn mod_div(a: usize, b: usize, p: usize) -> usize {
    a * mod_pow(b, p-2, p) % p
}

/// 二分探索
///
/// # Arguments
///
/// * `[left, right]`: 探索範囲
/// * `is_valid(i)` - iが有効かどうか
///
/// # Returns
/// * `Some(i)` - is_valid(i)がtrueになる最後のi
/// * `None` - is_valid(i)がすべてfalseになる場合
///
pub fn b_search_find_last<P>(left: usize, right: usize, is_valid: P) -> Option<usize> where
	P: Fn(usize) -> bool
{
	let mut min = left;
	let mut max = right;

	if min > max {
		None
	} else {
		while min < max {
			let c = max - (max - min) / 2;

			if is_valid(c) {
				min = c;
			} else {
				max = c - 1;
			}
		}

		if is_valid(min) {
			Some(min)
		} else {
			None
		}
	}
}

Submission Info

Judge Result