Take the 2-minute tour ×
Code Review Stack Exchange is a question and answer site for peer programmer code reviews. It's 100% free, no registration required.

I did an implementation of the Tonelli-Shanks algorithm as defined on Wikipedia. I put it here for review and sharing purpose.

Legendre Symbol implementation:

def legendre_symbol(a, p):
    """
    Legendre symbol
    Define if a is a quadratic residue modulo odd prime
    http://en.wikipedia.org/wiki/Legendre_symbol
    """
    ls = pow(a, (p - 1)/2, p)
    if ls == p - 1:
        return -1
    return ls

Prime modular square root (I just renamed the solution variable R to x and n to a):

def prime_mod_sqrt(a, p):
    """
    Square root modulo prime number
    Solve the equation
        x^2 = a mod p
    and return list of x solution
    http://en.wikipedia.org/wiki/Tonelli-Shanks_algorithm
    """
    a %= p

    # Simple case
    if a == 0:
        return [0]
    if p == 2:
        return [a]

    # Check solution existence on odd prime
    if legendre_symbol(a, p) != 1:
        return []

    # Simple case
    if p % 4 == 3:
        x = pow(a, (p + 1)/4, p)
        return [x, p-x]

    # Factor p-1 on the form q * 2^s (with Q odd)
    q, s = p - 1, 0
    while q % 2 == 0:
        s += 1
        q //= 2

    # Select a z which is a quadratic non resudue modulo p
    z = 1
    while legendre_symbol(z, p) != -1:
        z += 1
    c = pow(z, q, p)

    # Search for a solution
    x = pow(a, (q + 1)/2, p)
    t = pow(a, q, p)
    m = s
    while t != 1:
        # Find the lowest i such that t^(2^i) = 1
        i, e = 0, 2
        for i in xrange(1, m):
            if pow(t, e, p) == 1:
                break
            e *= 2

        # Update next value to iterate
        b = pow(c, 2**(m - i - 1), p)
        x = (x * b) % p
        t = (t * b * b) % p
        c = (b * b) % p
        m = i

    return [x, p-x]

If you have any optimization or found any error, please report it.

share|improve this question
 
There was an error in the final loop (c value was not properly set). This is now fixed. –  Phong Mar 2 at 13:50
3  
Hi Phong. On CodeReview we expect your code to have problems. When a reviewer points out a problem it makes the review look wrong if you update your question to be 'fixed'. Have a look at this meta post about what to do when you want to update your code after a review –  rolfl Mar 3 at 15:40
 
Thank you for the information. I will keep this in mind on my future post. –  Phong Mar 4 at 8:42
add comment

1 Answer

up vote 5 down vote accepted

Good job! I don't have a lot to comment on in this code. You have written straightforward clear code whose only complexity stems directly from the complexity of the operation it is performing. It would be good to include some of your external commentary (such as the renames of R and n) in the code itself to make it easier for someone to follow the documentation on wikipedia. You may want to include some of that documentation as well.

For reference, the rest of this review assumes that the code functions correctly; I don't have my math head on tonight.

There appears to be one case of redundant code, unless m can ever be 1, resulting in an empty range and thus no reassignment of i. Otherwise you can skip the assignment to i in the following:

i, e = 0, 2
for i in xrange(1, m):
    ...

There are a number of small strength-reduction optimizations you might consider, but in Python their impact is likely to be minimized - definitely profile before heading too deeply down the optimization path. For example in the following while loop:

# Factor p-1 on the form q * 2^s (with Q odd)
q, s = p - 1, 0
while q % 2 == 0:
    s += 1
    q //= 2

Both operations on q can be reduced. The modulus can be rewritten as a binary and q & 1, and the division as a binary shift q >>= 1. Alternately, you can use divmod to perform both operations at once.

Similarly, 2**(m - i - 1) is identical to 1 << (m - i - 1) for non-negative exponents.

share|improve this answer
 
Thanks for the review, this is a really nice feedback. I will remove i assignment (I always though i was only defined inside the for loop). I will also change the 2**(m - i - 1). I prefer let other micro-optimization unchanged in order to have better code readability. –  Phong Mar 3 at 15:25
 
Btw, i did not see how you would use divmod in the code –  Phong Mar 3 at 15:28
 
It would look something like while True: next_q, q_odd = divmod(q, 2); if q_odd: break; s += 1; q = next_q. The tricky question is whether avoiding the separate divisions would pay for the extra operations in Python, not to mention the tradeoffs in terms of clarity. –  Michael Urman Mar 4 at 3:18
add comment

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.