← run

algo-numalgo-01

0.615
8/13 tests· algorithms
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
```python
import math

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."""
    
    # Handle edge cases
    if m == 1:
        return 0
    if b % m == 1:
        return 0
    if a % m == 0 and b % m != 0:
        return -1
    if a % m == 0:
        return 0
    
    # Reduce a, b, and m modulo m
    a = a % m
    b = b % m
    
    # Handle case where a is not invertible modulo m
    if math.gcd(a, m) > 1:
        # Try to find a solution in the pre-period
        for x in range(1, m):
            if pow(a, x, m) == b:
                return x
        # If no solution in pre-period, return -1
        return -1
    
    # Use the baby-step giant-step algorithm
    m_root = math.isqrt(m) + 1
    table = {pow(a, i, m): i for i in range(m_root)}
    
    # Compute a^(-m_root) mod m
    a_inv = pow(a, m - 1 - m_root, m)
    
    # Compute b * a^(-m_root) mod m
    c = (b * pow(a_inv, m_root, m)) % m
    
    # Search for the solution in the table
    for r, i in table.items():
        if r == c:
            return i + m_root
    return -1
```

Test output (stdout)

.F....F.FFF..                                                            [100%]
=================================== FAILURES ===================================
_____________________________ test_simple_coprime ______________________________

    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)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E       assert -1 == 4
E        +  where -1 = discrete_log(3, 13, 17)

test_discrete_log.py:30: AssertionError
________________________________ test_base_zero ________________________________

    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
E       assert 0 == 1
E        +  where 0 = discrete_log(0, 0, 7)

test_discrete_log.py:66: AssertionError
_____________________ test_exhaustive_matches_brute_force ______________________

    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]
E       AssertionError: [(0, 0, 2, 0, 1), (0, 0, 3, 0, 1), (2, 2, 3, 3, 1), (0, 0, 4, 0, 1), (3, 3, 4, 4, 1), (0, 0, 5, 0, 1), ...]
E       assert not [(0, 0, 2, 0, 1), (0, 0, 3, 0, 1), (2, 2, 3, 3, 1), (0, 0, 4, 0, 1), (3, 3, 4, 4, 1), (0, 0, 5, 0, 1), ...]

test_discrete_log.py:87: AssertionError
________________ test_large_solvable_validity_and_minimal_bound ________________

    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
E           assert -1 != -1

test_discrete_log.py:99: AssertionError
_____________________ test_large_forces_subquadratic_prime _____________________

    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
E       assert 510471702489 == 196695866993
E        +  where 510471702489 = pow(7, 1654309, 999999999989)

test_discrete_log.py:111: AssertionError
=========================== short test summary info ============================
FAILED test_discrete_log.py::test_simple_coprime - assert -1 == 4
FAILED test_discrete_log.py::test_base_zero - assert 0 == 1
FAILED test_discrete_log.py::test_exhaustive_matches_brute_force - AssertionE...
FAILED test_discrete_log.py::test_large_solvable_validity_and_minimal_bound
FAILED test_discrete_log.py::test_large_forces_subquadratic_prime - assert 51...
5 failed, 8 passed in 1.62s
Self-repair attempts
[
  {
    "total": 13,
    "answer": "```python\nimport math\n\ndef discrete_log(a: int, b: int, m: int) -> int:\n    \"\"\"Return the smallest non-negative integer x such that\n    a**x is congruent to b (mod m)\n    or -1 if no such x exists.\"\"\"\n    \n    # Handle edge cases\n    if m == 1:\n        return 0\n    if b % m == 1:\n        return 0\n    if a % m == 0 and b % m != 0:\n        return -1\n    if a % m == 0:\n        return 0\n    \n    # Reduce a, b, and m modulo m\n    a = a % m\n    b = b % m\n    \n    # Handle case where a is not invertible modulo m\n    if math.gcd(a, m) > 1:\n        # Try to find a solution in the pre-period\n        for x in range(1, m):\n            if pow(a, x, m) == b:\n                return x\n        # If no solution in pre-period, return -1\n        return -1\n    \n    # Use the baby-step giant-step algorithm\n    m_root = math.isqrt(m) + 1\n    table = {pow(a, i, m): i for i in range(m_root)}\n    \n    # Compute a^(-m_root) mod m\n    a_inv = pow(a, m - 1 - m_root, m)\n    \n    # Compute b * a^(-m_root) mod m\n    c = (b * pow(a_inv, m_root, m)) % m\n    \n    # Search for the solution in the table\n    for r, i in table.items():\n        if r == c:\n            return i + m_root\n    return -1\n```",
    "passed": 8,
    "test_error": "        assert discrete_log(0, 1, 7) == 0\n>       assert discrete_log(0, 0, 7) == 1\nE       assert 0 == 1\nE        +  where 0 = discrete_log(0, 0, 7)\n\ntest_discrete_log.py:66: AssertionError\n_____________________ test_exhaustive_matches_brute_force ______________________\n\n    def test_exhaustive_matches_brute_force():\n        fails = []\n        for m in range(1, 70):\n            cap = m + 5  # covers the whole pre-period + period\n            for a in range(m):\n                for b in range(m):\n                    got = discrete_log(a, b, m)\n                    exp = brute(a, b, m, cap)\n                    if got != exp:\n                        fails.append((a, b, m, got, exp))\n>       assert not fails, fails[:10]\nE       AssertionError: [(0, 0, 2, 0, 1), (0, 0, 3, 0, 1), (2, 2, 3, 3, 1), (0, 0, 4, 0, 1), (3, 3, 4, 4, 1), (0, 0, 5, 0, 1), ...]\nE       assert not [(0, 0, 2, 0, 1), (0, 0, 3, 0, 1), (2, 2, 3, 3, 1), (0, 0, 4, 0, 1), (3, 3, 4, 4, 1), (0, 0, 5, 0, 1), ...]\n\ntest_discrete_log.py:87: AssertionError\n________________ test_large_solvable_validity_and_minimal_bound ________________\n\n    def test_large_solvable_validity_and_minimal_bound():\n        rng = random.Random(12345)\n        for _ in range(300):\n            m = rng.randint(2, 10 ** 9)\n            a = rng.randint(0, m - 1)\n            x = rng.randint(0, 2 * 10 ** 6)\n            b = pow(a, x, m)\n            r = discrete_log(a, b, m)\n>           assert r != -1\nE           assert -1 != -1\n\ntest_discrete_log.py:99: AssertionError\n_____________________ test_large_forces_subquadratic_prime _____________________\n\n    def test_large_forces_subquadratic_prime():\n        # ~1e12 prime; an O(m) scan would blow the time limit, an O(sqrt m) one won't\n        m = 999999999989\n        a = 7\n        x = 654321\n        b = pow(a, x, m)\n        r = discrete_log(a, b, m)\n>       assert pow(a, r, m) == b\nE       assert 510471702489 == 196695866993\nE        +  where 510471702489 = pow(7, 1654309, 999999999989)\n\ntest_discrete_log.py:111: AssertionError\n=========================== short test summary info ============================\nFAILED test_discrete_log.py::test_simple_coprime - assert -1 == 4\nFAILED test_discrete_log.py::test_base_zero - assert 0 == 1\nFAILED test_discrete_log.py::test_exhaustive_matches_brute_force - AssertionE...\nFAILED test_discrete_log.py::test_large_solvable_validity_and_minimal_bound\nFAILED test_discrete_log.py::test_large_forces_subquadratic_prime - assert 51...\n5 failed, 8 passed in 1.62s"
  },
  {
    "total": 13,
    "answer": "```python\nimport math\n\ndef discrete_log(a: int, b: int, m: int) -> int:\n    \"\"\"Return the smallest non-negative integer x such that\n    a**x is congruent to b (mod m)\n    or -1 if no such x exists.\"\"\"\n    \n    # Handle edge cases\n    if m == 1:\n        return 0\n    if b % m == 1:\n        return 0\n    if a % m == 0 and b % m != 0:\n        return -1\n    if a % m == 0:\n        return 0\n    \n    # Reduce a, b, and m modulo m\n    a = a % m\n    b = b % m\n    \n    # Handle case where a is not invertible modulo m\n    if math.gcd(a, m) > 1:\n        # Try to find a solution in the pre-period\n        for x in range(1, m):\n            if pow(a, x, m) == b:\n                return x\n        # If no solution in pre-period, return -1\n        return -1\n    \n    # Use the baby-step giant-step algorithm\n    m_root = math.isqrt(m) + 1\n    table = {pow(a, i, m): i for i in range(m_root)}\n    \n    # Compute a^(-m_root) mod m\n    a_inv = pow(a, m - 1 - m_root, m)\n    \n    # Compute b * a^(-m_root) mod m\n    c = (b * pow(a_inv, m_root, m)) % m\n    \n    # Search for the solution in the table\n    for r, i in table.items():\n        if r == c:\n            return i + m_root\n    return -1\n```",
    "passed": 8,
    "test_error": "        assert discrete_log(0, 1, 7) == 0\n>       assert discrete_log(0, 0, 7) == 1\nE       assert 0 == 1\nE        +  where 0 = discrete_log(0, 0, 7)\n\ntest_discrete_log.py:66: AssertionError\n_____________________ test_exhaustive_matches_brute_force ______________________\n\n    def test_exhaustive_matches_brute_force():\n        fails = []\n        for m in range(1, 70):\n            cap = m + 5  # covers the whole pre-period + period\n            for a in range(m):\n                for b in range(m):\n                    got = discrete_log(a, b, m)\n                    exp = brute(a, b, m, cap)\n                    if got != exp:\n                        fails.append((a, b, m, got, exp))\n>       assert not fails, fails[:10]\nE       AssertionError: [(0, 0, 2, 0, 1), (0, 0, 3, 0, 1), (2, 2, 3, 3, 1), (0, 0, 4, 0, 1), (3, 3, 4, 4, 1), (0, 0, 5, 0, 1), ...]\nE       assert not [(0, 0, 2, 0, 1), (0, 0, 3, 0, 1), (2, 2, 3, 3, 1), (0, 0, 4, 0, 1), (3, 3, 4, 4, 1), (0, 0, 5, 0, 1), ...]\n\ntest_discrete_log.py:87: AssertionError\n________________ test_large_solvable_validity_and_minimal_bound ________________\n\n    def test_large_solvable_validity_and_minimal_bound():\n        rng = random.Random(12345)\n        for _ in range(300):\n            m = rng.randint(2, 10 ** 9)\n            a = rng.randint(0, m - 1)\n            x = rng.randint(0, 2 * 10 ** 6)\n            b = pow(a, x, m)\n            r = discrete_log(a, b, m)\n>           assert r != -1\nE           assert -1 != -1\n\ntest_discrete_log.py:99: AssertionError\n_____________________ test_large_forces_subquadratic_prime _____________________\n\n    def test_large_forces_subquadratic_prime():\n        # ~1e12 prime; an O(m) scan would blow the time limit, an O(sqrt m) one won't\n        m = 999999999989\n        a = 7\n        x = 654321\n        b = pow(a, x, m)\n        r = discrete_log(a, b, m)\n>       assert pow(a, r, m) == b\nE       assert 510471702489 == 196695866993\nE        +  where 510471702489 = pow(7, 1654309, 999999999989)\n\ntest_discrete_log.py:111: AssertionError\n=========================== short test summary info ============================\nFAILED test_discrete_log.py::test_simple_coprime - assert -1 == 4\nFAILED test_discrete_log.py::test_base_zero - assert 0 == 1\nFAILED test_discrete_log.py::test_exhaustive_matches_brute_force - AssertionE...\nFAILED test_discrete_log.py::test_large_solvable_validity_and_minimal_bound\nFAILED test_discrete_log.py::test_large_forces_subquadratic_prime - assert 51...\n5 failed, 8 passed in 1.62s"
  },
  {
    "total": 13,
    "answer": "```python\nimport math\n\ndef discrete_log(a: int, b: int, m: int) -> int:\n    \"\"\"Return the smallest non-negative integer x such that\n    a**x is congruent to b (mod m)\n    or -1 if no such x exists.\"\"\"\n    \n    # Handle edge cases\n    if m == 1:\n        return 0\n    if b % m == 1:\n        return 0\n    if a % m == 0 and b % m != 0:\n        return -1\n    if a % m == 0:\n        return 0\n    \n    # Reduce a, b, and m modulo m\n    a = a % m\n    b = b % m\n    \n    # Handle case where a is not invertible modulo m\n    if math.gcd(a, m) > 1:\n        # Try to find a solution in the pre-period\n        for x in range(1, m):\n            if pow(a, x, m) == b:\n                return x\n        # If no solution in pre-period, return -1\n        return -1\n    \n    # Use the baby-step giant-step algorithm\n    m_root = math.isqrt(m) + 1\n    table = {pow(a, i, m): i for i in range(m_root)}\n    \n    # Compute a^(-m_root) mod m\n    a_inv = pow(a, m - 1 - m_root, m)\n    \n    # Compute b * a^(-m_root) mod m\n    c = (b * pow(a_inv, m_root, m)) % m\n    \n    # Search for the solution in the table\n    for r, i in table.items():\n        if r == c:\n            return i + m_root\n    return -1\n```",
    "passed": 8,
    "test_error": ""
  }
]