Python Code Review - Project Euler #12

Question

I just figured out Project Euler problem #12. This the first one I've done in Python, so I'm putting my code out there to hopefully get some constructive criticism.

If anyone is interested, would you mind looking over my solution and letting me know how I can improve on it, either through optimization or style?

Thanks!

from math import*

def divisorCount(n):
    nDivisors = 0

    integerSqrt = int(sqrt(n))

    if sqrt(n) == integerSqrt:
        nDivisors += 1

    for i in range(1, integerSqrt - 1):
        if n % i == 0:
            nDivisors += 2

    return nDivisors


def main():

    divisors = 0

    incrementTriangle = 1
    triangleNumber = 1

    while True:

        divisors = divisorCount(triangleNumber)

        if divisors < 500:
            incrementTriangle += 1
            triangleNumber = triangleNumber + incrementTriangle
        elif divisors >= 500:
            print "The first triangle number w/ 500+ divisors is", triangleNumber
            return 0


main()

Answer 1

Cygal's answer covers style, but you also asked for optimisation so here goes.

Your divisorCount can be a lot more efficient. The number of divisors can be calculated cia the following formula:

d(n) = prod((k + 1)) for a^k in the prime factorisation of n

For example, consider the prime factorisation of 360:

360 = 2^3 * 3^2 * 5^1

The exponents in the prime factorisation are 3, 2, and 1, so:

d(360) = (3 + 1) * (2 + 1) * (1 + 1)
       = 24

Note that 360 indeed has 24 factors: [1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 18, 20, 24, 30, 36, 40, 45, 60, 72, 90, 120, 180, 360].

Combine this with the fact that a number n can have at most one prime factor greater than sqrt(n), and you get something like the following:

# Calculates the number of positive integer divisors of n, including 1 and n
def d(n):
    divisors = 1
    stop = math.sqrt(n)

    primes = # Calculate prime numbers up to sqrt(n) in ascending order
    for prime in primes:
        if prime > stop:
            break

        exponent = 0        
        while n % prime == 0:
            n /= prime
            exponent += 1

        if a != 0:
            divisors *= (a + 1)
            stop = max(stop, n)

    # At this point, n will either be 1 or the only prime factor of n > root
    # in the later case, we have (n^1) as a factor so multiply divisors by (1 + 1)
    if n > root:
        divisors *= 2

    return divisors

Since you are going to be calling this function many times consecutively, it makes sense to pre-calculate a list of primes and pass it to the divisor function rather than calculating it each time. One simple and efficient method is the Sieve of Eratosthanes which can be implemented very easily:

# Compute all prime numbers less than k using the Sieve of Eratosthanes
def sieve(k):
    s = set(range(3, k, 2))
    s.add(2)

    for i in range(3, k, 2):
        if i in s:
            for j in range(i ** 2, k, i * 2):
                s.discard(j)

    return sorted(s)

This function can also fairly easily be modified to compute a range in pieces, e.g. to calculate the primes to 1M and then to 2M without recalculating the first 1M.

Putting this together your main loop will then look something like this:

primes = [2, 3, 5]
triangularNumber = 0
increment = 1
maxDivisors = 0
while maxDivisors < 500:
    triangularNumber += increment
    increment += 1

    if math.sqrt(triangularNumber) > primes[-1]:
       primes = sieve(primes, primes[-1] * 4)

    divisors = d(triangularNumber, primes)
    maxDivisors = max(divisors, maxDivisors)

print triangularNumber, "has", maxDivisors, "divisors."

As far as performance goes, here are the triangular numbers with the largest number of divisors and the times (in seconds) taken to get there (py1 is your code, py2 is using primes, c++ is the same as py2 but in c++):

TriNum      Value     d(n)  py1  py2  c++
 12375      76576500   576    0    0    0
 14399     103672800   648    0    0    0
 21735     236215980   768    1    0    0
 41040     842161320  1024    4    0    0
 78624    3090906000  1280  128   13    9
 98175    4819214400  1344  350   23   15
123200    7589181600  1512  702   38   23
126224    7966312200  1536  749   41   24
165375   13674528000  1728   -    74   42
201824   20366564400  1920   -   107   61
313599   49172323200  2304   -   633  153
395199   78091322400  2880   -   790  259
453375  102774672000  3072   -   910  358

There are probably other optimisations you can make based on the properties of triangular numbers, but I don't know the math. :)

Answer 2

from math import*

Missing space?

    for i in range(1, integerSqrt - 1):
        if n % i == 0:
            nDivisors += 2

I'm not saying it's better here, but you can achieve a lot with list comprehensions. For example, 2 * len([i for i in range(1, integerSqrt -1) if n % i == 0] would be equivalent. And less readable, so don't use it here!

            incrementTriangle += 1
            triangleNumber = triangleNumber + incrementTriangle

Be consistent with increments (triangleCount += incrementTriangle)

        elif divisors >= 500:
            print "The first triangle number w/ 500+ divisors is", triangleNumber
            return 0

Using returns or breaks in a loop is usually a bad idea if it can be avoided. It means that you need to scan the whole loop to see where the loop could end earlier than expected. In this case, what's preventing you from doing while divisors < 500. As an added bonus, you would stop mixing output and actual logic.

print "The first triangle number w/ 500+ divisors is", triangleNumber

Use the print() function provided in Python 2.6. It will make your code Python 3 compatible and.. more pythonic. If you want more flexibly, also consider using format strings. This would give:

print("The first triangle number w/ 500+ divisors is %d" % triangleNumber)

So, you have a main function and call it like this:

main()

You should instead use the __main__ idiom, as it will allow your file to be imported (for example to use divisorsCount).

And last, but not least, you should really follow the PEP 8, as others said in the comments. Consistency is really important to make your code more readable!