Solving double squares problem posted by Facebook

fbhack01.py

'''
Created on Jan 12, 2011

@author: cobpez
'''
import math

def double_squares(v):
    for x in xrange(int(math.sqrt(v))):
        y = math.sqrt(v - x*x)
        if not y % 1:
            yield x, y
        if x > y:
            break

INPUT="""20
1048039120
325
4
1475149141
1991891221
1816371419
4005625
2147483645
1740798996
138125
510644794
77068225
625058908
71825
2147483643
713302969
473200074
1
2147483642
2082925"""

_INPUT="""2
25
325
1
10"""


lines = [int(s) for s in INPUT.split()]
num_lines = len(lines)

def run():
    i = 0
    while i < len(lines):
        cases = lines[i]
        for j in range(cases):
            v = lines[i + j + 1]
            yield len(list(double_squares(v)))
        i += cases + 1

print list(run())

from timeit import Timer
time = Timer('run()', 'from __main__ import run')
RUNS = 1000000
REPEATS = 10
print sum(time.repeat(REPEATS, RUNS))/REPEATS