"Jedi's riddle" challenge

Question

The challenge can be found here. In short we have a recursive function, and we need to find when \$z_i\$ is equal to one.

$$ \begin{align} z_1 &= 1\\ z_i &= z_{i-1} * a\ \%\ n\\ \text{find}\ z_i &= 1 \end{align} $$

Input range:

$$ 2 \le a \lt n \le 10^9 $$

I have written the following. It is simple, but too slow.

import math

def solution(a, n):
    if math.gcd(a, n) > 1:
        return 0
    z = 1
    i = 1
    while True:
        z = z * a % n
        if z == 1:
            return i
        i += 1
l = map(int, input().split(' '))
print(solution(*l))

Answer 1

Code Review

Your code is a little hard to read.

You should have a blank line after the solution body, to separate it from the mainline code.

This code is hard to describe, document, and debug:

l = map(int, input().split(' '))
print(solution(*l))

What is l? How many arguments are there? If the wrong number of arguments are given as input, the problem doesn't occur immediately in the first statement but rather when attempting to execute the second.

It would be significantly clearer to write:

a, n = map(int, input().split())
print(solution(a, n))

Now we can see we expect two arguments to be given, and we have names for the arguments. If too many or too few arguments are given, the first line will raise the exception, which should mean it will be easier to debug.

Note: I'm using .split() instead of .split(' ') so that it will split the string on any white space character, including tabs, and multiple white-space characters will be considered a single delimiter. It is slightly friendlier.

Adding type hints and a docstring would be useful. Also, a main guard should be used.

Improved code:

import math

def solution(a: int, n: int) -> int:
    """
    Jedi Challenge:

    Return the smallest positive integer such that
        a^i = 1 (mod n)
    Returns 0 if there is no such number

    Constraints:
        2 <= a < n <= 10^9
    """

    if math.gcd(a, n) > 1:
        return 0

    z = 1
    i = 1
    while True:
        z = z * a % n
        if z == 1:
            return i
        i += 1

if __name__ == '__main__':
    a, n = map(int, input().split())
    print(solution(a, n))

Performance

This is not really a programming challenge; it is a math challenge.

First, due to the modulo operation, z = z * a % n will only ever produce a value between 0 and n-1. This means there are only n possible results at each iteration step. However, if 0 is produced, each iteration step afterwards will only ever produce 0, leaving n-1 non-trivial results.

If the k-th iteration produces 1, the code yields the result k. But if iteration continued, since it has reached its initial starting point, it would produce the same sequence of values in a never ending cycle.

If you chose a different starting value, that did not appear in the above cycle, it would must have one of 3 fates. It would either terminate at zero, create its own loop of unique values, or possibly merge into the existing sequence.

You have already determined that if gcd(a, n) != 1 then the iteration sequence would eventually reach zero, without finding a solution.

In fact, if n is prime, then \$a^{n-1} \equiv 1 \mod n\$, and therefore the iteration will produce the 1 after n-1 iterations, and must have produced all of the possible values between 1 and n-1.

So, as one optimization, if you can determine n is prime, you could immediately output the answer n-1. If not, then you could revert to doing the iterative calculations.

If n is not prime, you might want to look at the prime factorization of n, and see if you can determine a relation between it and the result. (Hint: there is.)