Naively counting Pythagorean triples in Python and Julia

JuliaGlobals

total = 0
N = 300

start_time = time()

for a in 0:(N - 1)
	for b in 0:(N - 1)
		for c in 0:(N - 1)
			if a^2 + b^2 == c^2
				total = total + 1
			end
		end
	end
end

end_time = time()

end_time - start_time

# Repeat now that JIT has done its work

total = 0
N = 300

start_time = time()

for a in 0:(N - 1)
	for b in 0:(N - 1)
		for c in 0:(N - 1)
			if a^2 + b^2 == c^2
				total = total + 1
			end
		end
	end
end

end_time = time()

println(end_time - start_time)

JuliaLocals

function loop(N::Integer)
	total = 0

	start_time = time()

	for a in 0:(N - 1)
		for b in 0:(N - 1)
			for c in 0:(N - 1)
				if a^2 + b^2 == c^2
					total = total + 1
				end
			end
		end
	end

	end_time = time()

	return end_time - start_time
end

loop(300)

# Repeat now that JIT has done its work

println(loop(300))

Python

total = 0
N = 300

from time import time

start = time()

for a in range(N):
	for b in range(N):
		for c in range(N):
			if a**2 + b**2 == c**2:
				total = total + 1

end = time()

print(end - start)

On my MacBook Pro, the timings are:

• Python - ~7000ms
• Julia w/ Globals - ~500ms
• Julia w/ Locals - ~10ms

The array version is intentionally very compressed, but still, I'd rather be able to use arrays to solve this sort of problem instead of three nested loops. Does Julia perform as well on that front?

def pythag_triples_loops(n):
    "Count Pythagorean triples using loops"
    total = 0
    for a in range(n):
        for b in range(n):
            for c in range(n):
                if a**2 + b**2 == c**2:
                    total = total + 1
    return total

def pythag_triples(n):
    "Count Pythagorean triples for an array 0, ..., n-1"
    x = np.arange(n)
    return np.sum(np.in1d(reduce(lambda x, y: x**2 + y**2, np.meshgrid(x, x)), x**2))

%timeit -n3  pythag_triples_loops(300)
#  3 loops, best of 3: 4.57 s per loop

%timeit pythag_triples(300)
#10 loops, best of 3: 17.5 ms per loop

Under PyPy I got a mean of 45ms over 100 runs, also on a recent Macbook Pro. 4.5x faster than the PyPy JIT is pretty awesome.

Using cython via ipython notebook:
%load_ext cythonmagic

paste this into the next cell block:

%%cython
import numpy as np

cimport numpy as np
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef calc(N=300):
    cdef:
        np.int total = 0
        Py_ssize_t i, j, k
    for i in xrange(N):
        for j in xrange(N):
            for k in xrange(N):
                if i**2 + j**2 == k**2:
                    total = total + 1
    return total

paste this into the next cell block:
%timeit -n3 calc()

3 loops, best of 3: 19.1 ms per loop

If N is defined as const in JuliaGlobals, the time changes from 660 ms to 34 ms on my machine