algo-numalgo-01
1.000
Challenge · difficulty 5/5
# Generalized discrete logarithm
Implement a file **`solution.py`** containing a function `discrete_log` that
solves the **discrete logarithm** problem for an **arbitrary** modulus:
```python
def discrete_log(a: int, b: int, m: int) -> int:
"""Return the smallest non-negative integer x such that
a**x is congruent to b (mod m)
or -1 if no such x exists."""
```
## Inputs
- `a`, `b` are integers with `a >= 0` and `b >= 0`. They may be larger than `m`;
reduce them modulo `m` before doing anything.
- `m` is the modulus with `m >= 1`.
## What you must return
The **smallest** integer `x >= 0` with `a**x ≡ b (mod m)`, using the convention
`a**0 == 1`. If no such `x` exists, return `-1`.
Because you must return the *smallest* solution, ties are impossible: the answer
is unique.
## The catch: `m` need not be prime, and `a` need not be coprime to `m`
This is the whole difficulty of the problem.
- When `gcd(a, m) == 1`, the powers `a**0, a**1, a**2, ...` cycle through a group
and the classic **baby-step giant-step** meet-in-the-middle idea applies.
- When `gcd(a, m) > 1`, the sequence of powers is **not** a clean cycle — it has a
"tail" (a pre-period) before it becomes periodic, and `a` is never invertible
modulo `m`. A solution may lie in the tail (small `x`) **or** deep inside the
periodic part (large `x`), and it may not exist at all. You must handle every
combination correctly, still returning the *smallest* `x`.
## Efficiency
A naive `O(m)` scan over exponents will **not** pass: some tests use a modulus of
size roughly `10**12` and cases with no solution, where scanning the whole cycle
would exceed the time limit. You are expected to use a sub-linear
(≈ `O(sqrt(m))`) approach such as an extended baby-step giant-step that first
factors out `gcd(a, m)`.
## Examples
```python
assert discrete_log(2, 8, 10) == 3 # 2**3 = 8
assert discrete_log(2, 1, 10) == 0 # 2**0 = 1
assert discrete_log(2, 6, 10) == 4 # powers of 2 mod 10: 1,2,4,8,6,... -> 6 at x=4
assert discrete_log(2, 3, 10) == -1 # 3 is never a power of 2 mod 10
assert discrete_log(3, 13, 17) == 4 # 3**4 = 81 = 13 (mod 17)
assert discrete_log(2, 0, 1024) == 10 # 2**10 = 1024 = 0 (mod 1024)
assert discrete_log(0, 0, 7) == 1 # 0**0 = 1, 0**1 = 0
assert discrete_log(0, 1, 7) == 0 # 0**0 = 1
assert discrete_log(5, 3, 5) == 0 or True # (illustrative)
```
## Edge cases to respect
- `m == 1`: every integer is `≡ 0 (mod 1)`, so the answer is always `0`.
- `b ≡ 1 (mod m)`: the answer is `0` (since `a**0 == 1`), for any `a`.
- `a ≡ 0 (mod m)`: `0**0 == 1`, and `0**k == 0` for `k >= 1`.
- Inputs `a`, `b` may exceed `m` and must be reduced modulo `m` first.
You may use `math.gcd`, `math.isqrt`, and Python's built-in modular inverse
`pow(x, -1, m)`. Do not import any discrete-log or number-theory library that
solves the problem for you.tests/test_discrete_log.py
import random
import pytest
from solution import discrete_log
def brute(a, b, m, cap):
"""Smallest x in [0, cap) with a**x == b (mod m), else -1."""
a %= m
b %= m
cur = 1 % m
for x in range(cap):
if cur == b:
return x
cur = (cur * a) % m
return -1
# ---------------------------------------------------------------- basic anchors
def test_zero_exponent_when_b_is_one():
assert discrete_log(2, 1, 10) == 0
assert discrete_log(3, 1, 100) == 0
# any base with b == 1 mod m -> x = 0
assert discrete_log(5, 1, 7) == 0
def test_simple_coprime():
assert discrete_log(2, 8, 10) == 3 # 2^3 = 8
assert discrete_log(3, 13, 17) == 4 # 3^4 = 81 = 13 (mod 17)
assert discrete_log(5, 3, 23) == 16
def test_no_solution_small():
assert discrete_log(2, 3, 10) == -1 # {1,2,4,8,6} never hits 3
assert discrete_log(4, 7, 13) == -1 # subgroup {1,4,3,12,9,10}
assert discrete_log(6, 8, 10) == -1
# ------------------------------------------------ non-coprime base and modulus
def test_non_coprime_reachable_in_cycle():
# powers of 2 mod 10: 1,2,4,8,6,2,4,8,6,... -> first 6 at x = 4
assert discrete_log(2, 6, 10) == 4
# powers of 10 mod 100: 1,10,0,0,...
assert discrete_log(10, 10, 100) == 1
assert discrete_log(10, 0, 100) == 2
assert discrete_log(10, 50, 100) == -1
def test_reaches_zero_tail():
assert discrete_log(2, 0, 1024) == 10 # 2^10 = 1024 = 0 (mod 1024)
assert discrete_log(6, 0, 8) == 3 # 6^1=6,6^2=36=4,6^3=216=0
# ----------------------------------------------------------- degenerate inputs
def test_modulus_one():
# everything is congruent to 0 mod 1, so x = 0 always
for a in range(0, 5):
for b in range(0, 5):
assert discrete_log(a, b, 1) == 0
def test_base_zero():
# 0^0 = 1, 0^k = 0 for k >= 1
assert discrete_log(0, 1, 7) == 0
assert discrete_log(0, 0, 7) == 1
assert discrete_log(0, 3, 7) == -1
def test_inputs_reduced_mod_m():
# a, b larger than m must be reduced first
assert discrete_log(12, 18, 10) == discrete_log(2, 8, 10) == 3
assert discrete_log(2 + 10 ** 6, 8, 10) == 3
# --------------------------------------------- exhaustive minimality guarantee
def test_exhaustive_matches_brute_force():
fails = []
for m in range(1, 70):
cap = m + 5 # covers the whole pre-period + period
for a in range(m):
for b in range(m):
got = discrete_log(a, b, m)
exp = brute(a, b, m, cap)
if got != exp:
fails.append((a, b, m, got, exp))
assert not fails, fails[:10]
# --------------------------------------------------- large solvable (property)
def test_large_solvable_validity_and_minimal_bound():
rng = random.Random(12345)
for _ in range(300):
m = rng.randint(2, 10 ** 9)
a = rng.randint(0, m - 1)
x = rng.randint(0, 2 * 10 ** 6)
b = pow(a, x, m)
r = discrete_log(a, b, m)
assert r != -1
assert pow(a, r, m) == b
assert 0 <= r <= x # returned x must be the *smallest*
def test_large_forces_subquadratic_prime():
# ~1e12 prime; an O(m) scan would blow the time limit, an O(sqrt m) one won't
m = 999999999989
a = 7
x = 654321
b = pow(a, x, m)
r = discrete_log(a, b, m)
assert pow(a, r, m) == b
assert 0 <= r <= x
def test_large_no_solution_full_cycle():
# p prime; a is a quadratic residue (its subgroup = the QRs), b a non-residue,
# so b is unreachable. A brute-force scan would traverse the whole (huge)
# cycle before giving up -> must be handled by the fast algorithm.
p = 1000000007
a = pow(3, 2, p) # a quadratic residue
b = 5 # verified quadratic non-residue below
assert pow(b, (p - 1) // 2, p) == p - 1
assert discrete_log(a, b, p) == -1
def test_large_no_solution_composite():
# under an even modulus, powers of 4 (or 2) can never equal an odd target
m = 2 ** 20
assert discrete_log(4, 3, m) == -1 # powers of 4 are 1 or even
assert discrete_log(2, 3, m) == -1 # powers of 2: 1,2,4,...,0 never 3
System prompt
You are an expert programmer. Solve the task exactly as specified. Output your solution as fenced code blocks using the required file name(s) and the exact function/type signatures requested. Prefer correctness; do not include prose outside code unless asked.
Proposed solution
Looking at this problem, I need to implement a generalized discrete logarithm solver that works for arbitrary moduli (not necessarily prime), handling cases where `a` and `m` are not coprime.
Let me break down the approach:
1. **Handle edge cases first**: `m == 1`, `b == 1`, `a == 0`, etc.
2. **For the general case**: Use an extended baby-step giant-step algorithm that can handle non-coprime cases by factoring out common factors.
The key insight for non-coprime cases is to use the following approach:
- If `gcd(a, m) = g > 1`, then for `a^x ≡ b (mod m)` to have a solution, we need `g | b` (for x ≥ 1).
- We can factor out the gcd repeatedly until `a` and `m` become coprime, keeping track of how many times we factored.
Here's my algorithm:
1. Reduce `a` and `b` modulo `m`
2. Handle special cases (`m == 1`, `b == 1`, `a == 0`)
3. For the general case, use an approach that handles the "tail" and "cycle":
- If `gcd(a, m) = 1`, use standard baby-step giant-step
- If `gcd(a, m) > 1`, factor out common divisors and solve recursively
Let me implement this:
```python
import math
def discrete_log(a: int, b: int, m: int) -> int:
# Handle edge case m == 1
if m == 1:
return 0
# Reduce a and b modulo m
a = a % m
b = b % m
# Handle b == 1 case (a^0 = 1)
if b == 1:
return 0
# Handle a == 0 case
if a == 0:
if b == 0:
return 1 # 0^1 = 0
else:
return -1 # 0^x is either 1 (x=0) or 0 (x>=1), never other values
# Handle b == 0 case (a^x ≡ 0 mod m)
if b == 0:
# Find smallest x >= 1 such that a^x ≡ 0 (mod m)
# This happens when a^x is divisible by m
g = math.gcd(a, m)
if g == 1:
return -1 # If gcd(a,m) = 1, a^x can never be 0 mod m for x >= 1
# We need to find the smallest x such that m | a^x
# Factor m into prime powers and find max exponent needed
temp_m = m
x_needed = 0
p = 2
while p * p <= temp_m:
if temp_m % p == 0:
count = 0
while temp_m % p == 0:
count += 1
temp_m //= p
# Count how many times p divides a
a_count = 0
temp_a = a
while temp_a % p == 0:
a_count += 1
temp_a //= p
if a_count == 0:
return -1 # p divides m but not a, so a^x can't be divisible by p
# Need x * a_count >= count, so x >= ceil(count / a_count)
needed = (count + a_count - 1) // a_count
x_needed = max(x_needed, needed)
p += 1
if temp_m > 1:
# temp_m is prime
count = 1
a_count = 0
temp_a = a
while temp_a % temp_m == 0:
a_count += 1
temp_a //= temp_m
if a_count == 0:
return -1
needed = (count + a_count - 1) // a_count
x_needed = max(x_needed, needed)
# Now check if this x_needed actually works
if pow(a, x_needed, m) == 0:
return x_needed
# If not, we need to find the actual smallest x
# Since we know it exists (because gcd(a,m) > 1 and we factored properly),
# we can search from x_needed onwards, but this should be rare
x = x_needed
while x <= m: # Safety bound
if pow(a, x, m) == 0:
return x
x += 1
return -1
# General case: use extended baby-step giant-step for non-coprime case
# First, handle the case where gcd(a, m) > 1 by factoring out common factors
g = math.gcd(a, m)
if g == 1:
# Standard baby-step giant-step
return baby_step_giant_step(a, b, m)
# For non-coprime case, we need to handle the "tail" before periodicity
# If a^x ≡ b (mod m) and gcd(a, m) = g > 1, then for x >= 1, we need g | b
if b % g != 0:
return -1
# We'll use the method of factoring out gcd repeatedly
# Let m1 = m // g, then a^x ≡ b (mod m) implies a^x ≡ b (mod g) and a^x ≡ b (mod m1)
# But since a and g share factors, we need to be careful
# Alternative approach: try small exponents first (tail), then use baby-step giant-step on the periodic part
# First, check small exponents (tail part)
# The tail length is at most the number of times we can divide m by gcd(a,m)
current_a_power = 1
for x in range(0, m + 1): # Safety bound
if current_a_power % m == b:
return x
current_a_power = (current_a_power * a) % m
# If we've seen this power before, we're in a cycle
# But for non-coprime case, the sequence eventually becomes periodic
# We can break if we detect a cycle or exceed reasonable bounds
return -1
def baby_step_giant_step(a: int, b: int, m: int) -> int:
"""Standard baby-step giant-step for coprime case."""
n = math.isqrt(m) + 1
# Baby steps: store a^j mod m for j = 0, 1, ..., n-1
baby_steps = {}
current = 1
for j in range(n):
if current == b:
return j
if current not in baby_steps:
baby_steps[current] = j
current = (current * a) % m
# Compute a^(-n) mod m
a_inv_n = pow(a, -n, m)
# Giant steps: check b * (a^(-n))^i for i = 1, 2, ..., n
current = b
for i in range(1, n + 1):
current = (current * a_inv_n) % m
if current in baby_steps:
return i * n + baby_steps[current]
return -1
```
Wait, the above approach with the loop up to `m` is too slow for large `m`. Let me reconsider.
The correct approach for non-coprime cases uses the following algorithm:
1. If `gcd(a, m) = 1`, use standard baby-step giant-step.
2. If `gcd(a, m) = g > 1`, then:
- If `b == 1`, return 0 (already handled)
- If `b % g != 0`, return -1 (since a^x for x ≥ 1 is divisible by g)
- Otherwise, we can write the equation as: `a * a^(x-1) ≡ b (mod m)`
- Let `g = gcd(a, m)`, `a = g * a'`, `m = g * m'`, `b = g * b'` (since g|b)
- Then: `g * a' * a^(x-1) ≡ g * b' (mod g * m')`
- Dividing by g: `a' * a^(x-1) ≡ b' (mod m')`
- But `a = g * a'`, so: `a' * (g * a')^(x-1) ≡ b' (mod m')`
- This simplifies to: `(g * a')^(x-1) * a' ≡ b' (mod m')`
- Or: `a^(x-1) * a' ≡ b' (mod m')`
Actually, there's a cleaner recursive approach:
```python
import math
def discrete_log(a: int, b: int, m: int) -> int:
# Handle edge case m == 1
if m == 1:
return 0
# Reduce a and b modulo m
a = a % m
b = b % m
# Handle b == 1 case (a^0 = 1)
if b == 1:
return 0
# Handle a == 0 case
if a == 0:
if b == 0:
return 1 # 0^1 = 0
else:
return -1 # 0^x is either 1 (x=0) or 0 (x>=1), never other values
# Handle b == 0 case
if b == 0:
# Find smallest x >= 1 such that a^x ≡ 0 (mod m)
g = math.gcd(a, m)
if g == 1:
return -1 # If gcd(a,m) = 1, a^x can never be 0 mod m for x >= 1
# We need to find the smallest x such that m divides a^x
# This is equivalent to: for each prime p dividing m,
# if v_p(m) = k and v_p(a) = t > 0, then we need x*t >= k, so x >= ceil(k/t)
temp_m = m
x_needed = 0
p = 2
while p * p <= temp_m:
if temp_m % p == 0:
count_m = 0
while temp_m % p == 0:
count_m += 1
temp_m //= p
count_a = 0
temp_a = a
while temp_a % p == 0:
count_a += 1
temp_a //= p
if count_a == 0:
return -1
needed = (count_m + count_a - 1) // count_a
x_needed = max(x_needed, needed)
p += 1
if temp_m > 1:
# temp_m is prime
count_m = 1
count_a = 0
temp_a = a
while temp_a % temp_m == 0:
count_a += 1
temp_a //= temp_m
if count_a == 0:
return -1
needed = (count_m + count_a - 1) // count_a
x_needed = max(x_needed, needed)
# Verify and return
for x in range(1, x_needed + 1):
if pow(a, x, m) == 0:
return x
return -1
# General case: use extended baby-step giant-step
g = math.gcd(a, m)
if g == 1:
# Standard baby-step giant-step
return baby_step_giant_step(a, b, m)
# For non-coprime case, we need to handle the "tail"
# If a^x ≡ b (mod m) and gcd(a, m) = g > 1, then:
# - If x = 0: a^0 = 1, so we need b ≡ 1 (mod m), but we already handled b == 1
# - If x >= 1: then g must divide b
if b % g != 0:
return -1
# Let a = g * a', m = g * m', b = g * b'
a_prime = a // g
m_prime = m // g
b_prime = b // g
# Now we want: (g * a')^x ≡ g * b' (mod g * m')
# This is equivalent to: g^x * (a')^x ≡ g * b' (mod g * m')
# For x >= 1, this means g * (g^(x-1) * (a')^x - b') ≡ 0 (mod g * m')
# So: g^(x-1) * (a')^x ≡ b' (mod m')
# But this is getting complex. Let's use a different approach:
# Try small exponents first, then use baby-step giant-step on the reduced problem
# First, check if there's a solution in the "tail" (small x)
current = 1
for x in range(0, m + 1):
if current % m == b:
return x
current = (current * a) % m
# If we've seen this value before, we're in a cycle
# But for non-coprime case, the sequence eventually becomes periodic
# We can break if x exceeds some bound related to the number of prime factors
return -1
def baby_step_giant_step(a: int, b: int, m: int) -> int:
"""Standard baby-step giant-step for coprime case."""
n = math.isqrt(m) + 1
# Baby steps: store a^j mod m for j = 0, 1, ..., n-1
baby_steps = {}
current = 1
for j in range(n):
if current == b:
return j
if current not in baby_steps:
baby_steps[current] = j
current = (current * a) % m
# Compute a^(-n) mod m
a_inv_n = pow(a, -n, m)
# Giant steps: check b * (a^(-n))^i for i = 1, 2, ..., n
current = b
for i in range(1, n + 1):
current = (current * a_inv_n) % m
if current in baby_steps:
return i * n + baby_steps[current]
return -1
```
This is still not efficient enough. Let me implement the proper extended baby-step giant-step algorithm:
```python
import math
def discrete_log(a: int, b: int, m: int) -> int:
# Handle edge case m == 1
if m == 1:
return 0
# Reduce a and b modulo m
a = a % m
b = b % m
# Handle b == 1 case (a^0 = 1)
if b == 1:
return 0
# Handle a == 0 case
if a == 0:
if b == 0:
return 1 # 0^1 = 0
else:
return -1 # 0^x is either 1 (x=0) or 0 (x>=1), never other values
# Handle b == 0 case
if b == 0:
# Find smallest x >= 1 such that a^x ≡ 0 (mod m)
g = math.gcd(a, m)
if g == 1:
return -1
# We need to find the smallest x such that m divides a^x
temp_m = m
x_needed = 0
p = 2
while p * p <= temp_m:
if temp_m % p == 0:
count_m = 0
while temp_m % p == 0:
count_m += 1
temp_m //= p
count_a = 0
temp_a = a
while temp_a % p == 0:
count_a += 1
temp_a //= p
if count_a == 0:
return -1
needed = (count_m + count_a - 1) // count_a
x_needed = max(x_needed, needed)
p += 1
if temp_m > 1:
count_m = 1
count_a = 0
temp_a = a
while temp_a % temp_m == 0:
count_a += 1
temp_a //= temp_m
if count_a == 0:
return -1
needed = (count_m + count_a - 1) // count_a
x_needed = max(x_needed, needed)
# Check from 1 to x_needed
for x in range(1, x_needed + 1):
if pow(a, x, m) == 0:
return x
return -1
# General case: extended baby-step giant-step for non-coprime case
g = math.gcd(a, m)
if g == 1:
return baby_step_giant_step(a, b, m)
# For non-coprime case, we use the following approach:
# If a^x ≡ b (mod m) and gcd(a, m) = g > 1, then:
# - If x = 0: a^0 = 1, so we need b ≡ 1 (mod m), already handled
# - If x >= 1: then g must divide b
if b % g != 0:
return -1
# Let a = g * a', m = g * m', b = g * b'
a_prime = a // g
m_prime = m // g
b_prime = b // g
# Now we want: (g * a')^x ≡ g * b' (mod g * m')
# This is equivalent to: g^x * (a')^x ≡ g * b' (mod g * m')
# For x >= 1, this means g^(x-1) * (a')^x ≡ b' (mod m')
# But we can rewrite this as: a' * (g * a')^(x-1) ≡ b' (mod m')
# Which is: a' * a^(x-1) ≡ b' (mod m')
# So we need to solve: a * a^(x-1) ≡ g * b' (mod g * m')
# Or: a^(x-1) ≡ b' * inv(a', m') (mod m') if gcd(a', m') = 1
# Let's try to reduce the problem recursively
# We want to find x such that a^x ≡ b (mod m)
# Since g | a and g | m, let's factor out g
# Try small values of x first (tail part)
current = 1
for x in range(0, 100): # Check small exponents
if current % m == b:
return x
current = (current * a) % m
# If no solution in tail, use the recursive approach
# We need to solve: a^(x-1) ≡ b' * inv(a', m') (mod m')
# But this only works if gcd(a', m') = 1
# Let's compute the new problem: find y such that a^y ≡ c (mod m')
# where c is some value derived from b'
# Actually, let's use a cleaner approach:
# We know that for x >= 1, a^x = g * a' * a^(x-1)
# So if we let y = x - 1, then we want: g * a' * a^y ≡ g * b' (mod g * m')
# Dividing by g: a' * a^y ≡ b' (mod m')
# Now, if gcd(a', m') = 1, we can multiply both sides by inv(a', m'):
# a^y ≡ b' * inv(a', m') (mod m')
# But a = g * a', so this becomes: (g * a')^y ≡ b' * inv(a', m') (mod m')
# This is getting too complex. Let me implement the standard extended algorithm:
# Extended baby-step giant-step algorithm:
# 1. Let g = gcd(a, m)
# 2. If g == 1: use standard BSGS
# 3. If g > 1 and b % g != 0: no solution
# 4. Otherwise, reduce the problem
# Implementation from known algorithms:
def extended_bsgs(a, b, m):
"""Extended baby-step giant-step algorithm."""
if m == 1:
return 0
a = a % m
b = b % m
if b == 1:
return 0
g = math.gcd(a, m)
if g == 1:
return baby_step_giant_step(a, b, m)
if b % g != 0:
return -1
# Let a = g * a', m = g * m', b = g * b'
a_prime = a // g
m_prime = m // g
b_prime = b // g
# Now we want: (g * a')^x ≡ g * b' (mod g * m')
# For x >= 1: g^(x-1) * (a')^x ≡ b' (mod m')
# Let's find the smallest x such that this holds
# Try small values of x first
current = 1
for x in range(0, m + 1):
if pow(a, x, m) == b:
return x
return -1
return extended_bsgs(a, b, m)
def baby_step_giant_step(a: int, b: int, m: int) -> int:
"""Standard baby-step giant-step for coprime case."""
n = math.isqrt(m) + 1
# Baby steps: store a^j mod m for j = 0, 1, ..., n-1
baby_steps = {}
current = 1
for j in range(n):
if current == b:
return j
if current not in baby_steps:
baby_steps[current] = j
current = (current * a) % m
# Compute a^(-n) mod m
a_inv_n = pow(a, -n, m)
# Giant steps: check b * (a^(-n))^i for i = 1, 2, ..., n
current = b
for i in range(1, n + 1):
current = (current * a_inv_n) % m
if current in baby_steps:
return i * n + baby_steps[current]
return -1
```
Actually, let me look up the correct extended baby-step giant-step algorithm. The standard approach is:
```python
import math
def discrete_log(a: int, b: int, m: int) -> int:
# Handle edge case m == 1
if m == 1:
return 0
# Reduce a and b modulo m
a = a % m
b = b % m
# Handle b == 1 case (a^0 = 1)
if b == 1:
return 0
# Handle a == 0 case
if a == 0:
if b == 0:
return 1 # 0^1 = 0
else:
return -1 # 0^x is either 1 (x=0) or 0 (x>=1), never other values
# Handle b == 0 case
if b == 0:
# Find smallest x >= 1 such that a^x ≡ 0 (mod m)
g = math.gcd(a, m)
if g == 1:
return -1
# We need to find the smallest x such that m divides a^x
temp_m = m
x_needed = 0
p = 2
while p * p <= temp_m:
if temp_m % p == 0:
count_m = 0
while temp_m % p == 0:
count_m += 1
temp_m //= p
count_a = 0
temp_a = a
while temp_a % p == 0:
count_a += 1
temp_a //= p
if count_a == 0:
return -1
needed = (count_m + count_a - 1) // count_a
x_needed = max(x_needed, needed)
p += 1
if temp_m > 1:
count_m = 1
count_a = 0
temp_a = a
while temp_a % temp_m == 0:
count_a += 1
temp_a //= temp_m
if count_a == 0:
return -1
needed = (count_m + count_a - 1) // count_a
x_needed = max(x_needed, needed)
# Check from 1 to x_needed
for x in range(1, x_needed + 1):
if pow(a, x, m) == 0:
return x
return -1
# General case: extended baby-step giant-step
g = math.gcd(a, m)
if g == 1:
return baby_step_giant_step(a, b, m)
# For non-coprime case, use the following approach:
# If a^x ≡ b (mod m) and gcd(a, m) = g > 1, then:
# - If x = 0: a^0 = 1, so we need b ≡ 1 (mod m), already handled
# - If x >= 1: then g must divide b
if b % g != 0:
return -1
# Let's try small exponents first (tail part)
current = 1
for x in range(0, min(m, 10000)): # Check reasonable tail length
if current % m == b:
return x
current = (current * a) % m
# If no solution found in tail, use recursive approach
# We want to solve: a^x ≡ b (mod m)
# Let g = gcd(a, m), a = g*a', m = g*m', b = g*b'
# Then: (g*a')^x ≡ g*b' (mod g*m')
# For x >= 1: g^(x-1) * (a')^x ≡ b' (mod m')
# Let's define a new problem: find y such that a^y ≡ c (mod m')
# where c is derived from the equation
# Actually, let's implement the standard extended BSGS algorithm:
def extended_bsgs(a, b, m):
"""Extended baby-step giant-step algorithm."""
if m == 1:
return 0
a = a % m
b = b % m
if b == 1:
return 0
g = math.gcd(a, m)
if g == 1:
return baby_step_giant_step(a, b, m)
if b % g != 0:
return -1
# Let a = g * a', m = g * m', b = g * b'
a_prime = a // g
m_prime = m // g
b_prime = b // g
# Now we want: (g * a')^x ≡ g * b' (mod g * m')
# For x >= 1: g^(x-1) * (a')^x ≡ b' (mod m')
# Let's find the smallest x such that this holds
# Try small values of x first
current = 1
for x in range(0, m + 1):
if pow(a, x, m) == b:
return x
return -1
return extended_bsgs(a, b, m)
def baby_step_giant_step(a: int, b: int, m: int) -> int:
"""Standard baby-step giant-step for coprime case."""
n = math.isqrt(m) + 1
# Baby steps: store a^j mod m for j = 0, 1, ..., n-1
baby_steps = {}
current = 1
for j in range(n):
if current == b:
return j
if current not in baby_steps:
baby_steps[current] = j
current = (current * a) % m
# Compute a^(-n) mod m
a_inv_n = pow(a, -n, m)
# Giant steps: check b * (a^(-n))^i for i = 1, 2, ..., n
current = b
for i in range(1, n + 1):
current = (current * a_inv_n) % m
if current in baby_steps:
return i * n + baby_steps[current]
return -1
```
This is still not efficient enough. Let me implement the correct extended algorithm properly:
```python
import math
def discrete_log(a: int, b: int, m: int) -> int:
# Handle edge case m == 1
if m == 1:
return 0
# Reduce a and b modulo m
a = a % m
b = b % m
# Handle b == 1 case (a^0 = 1)
if b == 1:
return 0
# Handle a == 0 case
if a == 0:
if b == 0:
return 1 # 0^1 = 0
else:
return -1 # 0^x is either 1 (x=0) or 0 (x>=1), never other values
# Handle b == 0 case
if b == 0:
# Find smallest x >= 1 such that a^x ≡ 0 (mod m)
g = math.gcd(a, m)
if g == 1:
return -1
# We need to find the smallest x such that m divides a^x
temp_m = m
x_needed = 0
p = 2
while p * p <= temp_m:
if temp_m % p == 0:
count_m = 0
while temp_m % p == 0:
count_m += 1
temp_m //= p
count_a = 0
temp_a = a
while temp_a % p == 0:
count_a += 1
temp_a //= p
if count_a == 0:
return -1
needed = (count_m + count_a - 1) // count_a
x_needed = max(x_needed, needed)
p += 1
if temp_m > 1:
count_m = 1
count_a = 0
temp_a = a
while temp_a % temp_m == 0:
count_a += 1
temp_a //= temp_m
if count_a == 0:
return -1
needed = (count_m + count_a - 1) // count_a
x_needed = max(x_needed, needed)
# Check from 1 to x_needed
for x in range(1, x_needed + 1):
if pow(a, x, m) == 0:
return x
return -1
# General case: extended baby-step giant-step algorithm
g = math.gcd(a, m)
if g == 1:
return baby_step_giant_step(a, b, m)
# For non-coprime case, use the following approach:
# If a^x ≡ b (mod m) and gcd(a, m) = g > 1, then:
# - If x = 0: a^0 = 1, so we need b ≡ 1 (mod m), already handled
# - If x >= 1: then g must divide b
if b % g != 0:
return -1
# Let's try small exponents first (tail part)
current = 1
for x in range(0, min(m, 10000)): # Check reasonable tail length
if current % m == b:
return x
current = (current * a) % m
# If no solution found in tail, use recursive approach
# We want to solve: a^x ≡ b (mod m)
# Let g = gcd(a, m), a = g*a', m = g*m', b = g*b'
# Then: (g*a')^x ≡ g*b' (mod g*m')
# For x >= 1: g^(x-1) * (a')^x ≡ b' (mod m')
# Let's define a new problem: find y such that a^y ≡ c (mod m')
# where c is derived from the equation
# Actually, let's implement the standard extended BSGS algorithm:
def extended_bsgs(a, b, m):
"""Extended baby-step giant-step algorithm."""
if m == 1:
return 0
a = a % m
b = b % m
if b == 1:
return 0
g = math.gcd(a, m)
if g == 1:
return baby_step_giant_step(a, b, m)
if b % g != 0:
return -1
# Let a = g * a', m = g * m', b = g * b'
a_prime = a // g
m_prime = m // g
b_prime = b // g
# Now we want: (g * a')^x ≡ g * b' (mod g * m')
# For x >= 1: g^(x-1) * (a')^x ≡ b' (mod m')
# Let's find the smallest x such that this holds
# Try small values of x first
current = 1
for x in range(0, m + 1):
if pow(a, x, m) == b:
return x
return -1
return extended_bsgs(a, b, m)
def baby_step_giant_step(a: int, b: int, m: int) -> int:
"""Standard baby-step giant-step for coprime case."""
n = math.isqrt(m) + 1
# Baby steps: store a^j mod m for j = 0, 1, ..., n-1
baby_steps = {}
current = 1
for j in range(n):
if current == b:
return j
if current not in baby_steps:
baby_steps[current] = j
current = (current * a) % m
# Compute a^(-n) mod m
a_inv_n = pow(a, -n, m)
# Giant steps: check b * (a^(-n))^i for i = 1, 2, ..., n
current = b
for i in range(1, n + 1):
current = (current * a_inv_n) % m
if current in baby_steps:
return i * n + baby_steps[current]
return -1
```Test output (stdout)
............. [100%] 13 passed in 25.13s