math/baby_step_giant_step.hpp

View this file on GitHub
Last update: 2024-02-23 10:38:03+09:00
Include: #include "math/baby_step_giant_step.hpp"

Depends on

Code

#pragma once

#include "math/inv_mod.hpp"
#include "math/pow_mod.hpp"

// find minimum K s.t. X ^ K = Y (mod M) and K >= lb
// O(sqrt(M))
long long baby_step_giant_step(long long X, long long Y, long long M, long long lb = 0LL) {
    X %= M, Y %= M;
    if (X < 0) X += M;
    if (Y < 0) Y += M;
    if (std::gcd(X, M) != 1) return -1;
    long long sqM = sqrt(M) + 1;
    std::unordered_map<long long, long long> mp;
    // Baby-step
    long long Z = pow_mod(X, lb, M);  // Z = X ^ lb (mod M)
    for (int i = 0; i < sqM; i++) {
        if (mp.find(Z) == mp.end()) mp[Z] = i + lb;
        Z = Z * X % M;
    }
    // Z = X ^ sqM (mod M)
    // R = Z ^ (-1) (mod M)
    long long R = inv_mod(pow_mod(X, sqM, M), M);
    // Giant-step
    for (int i = 0; i <= sqM; i++) {
        if (mp.find(Y) != mp.end()) return mp[Y] + i * sqM;
        Y = Y * R % M;
    }
    return -1;
}

#line 2 "math/baby_step_giant_step.hpp"

#line 2 "math/inv_mod.hpp"

#include <cassert>
#line 2 "math/extended_gcd.hpp"

#include <tuple>
// find (x, y) s.t. ax + by = gcd(a, b)
// a, b >= 0
// return {x, y, gcd(a, b)}
template <class T> std::tuple<T, T, T> extended_gcd(T a, T b) {
    if (b == 0) return {1, 0, a};
    auto [y, x, g] = extended_gcd(b, a % b);
    return {x, y - (a / b) * x, g};
}
#line 5 "math/inv_mod.hpp"

// mod >= 2
long long inv_mod(long long a, long long mod) {
    a %= mod;
    if (a < 0) a += mod;
    assert(a != 0);
    auto [x, y, g] = extended_gcd(a, mod);
    x %= mod;
    if (x < 0) x += mod;
    return x;
}
#line 2 "math/pow_mod.hpp"

long long pow_mod(long long a, long long n, const long long mod) {
    assert(n >= 0 and mod >= 1);
    if (mod == 1) return 0;
    a %= mod;
    if (a < 0) a += mod;
    long long res = 1;
    while (n) {
        if (n & 1) res = res * a % mod;
        a = a * a % mod;
        n >>= 1;
    }
    return res;
}
#line 5 "math/baby_step_giant_step.hpp"

// find minimum K s.t. X ^ K = Y (mod M) and K >= lb
// O(sqrt(M))
long long baby_step_giant_step(long long X, long long Y, long long M, long long lb = 0LL) {
    X %= M, Y %= M;
    if (X < 0) X += M;
    if (Y < 0) Y += M;
    if (std::gcd(X, M) != 1) return -1;
    long long sqM = sqrt(M) + 1;
    std::unordered_map<long long, long long> mp;
    // Baby-step
    long long Z = pow_mod(X, lb, M);  // Z = X ^ lb (mod M)
    for (int i = 0; i < sqM; i++) {
        if (mp.find(Z) == mp.end()) mp[Z] = i + lb;
        Z = Z * X % M;
    }
    // Z = X ^ sqM (mod M)
    // R = Z ^ (-1) (mod M)
    long long R = inv_mod(pow_mod(X, sqM, M), M);
    // Giant-step
    for (int i = 0; i <= sqM; i++) {
        if (mp.find(Y) != mp.end()) return mp[Y] + i * sqM;
        Y = Y * R % M;
    }
    return -1;
}