The problem statement:
An arithmetic progression is a sequence of the form a, a+b, a+2b, ..., a+nb where n=0, 1, 2, 3, ... . For this problem, a is a non-negative integer and b is a positive integer.
Write a program that finds all arithmetic progressions of length n in the set S of bisquares. The set of bisquares is defined as the set of all integers of the form p2 + q2 (where p and q are non-negative integers).
The two lines of input are n and m, which are the length of each sequence, and the upper bound to limit the search of the bi squares respectively.
My code solves the problem correctly, however it is too slow for the 5 second time constraint. Is there any way to further optimize my code?
import heapq
with open('ariprog.in', 'r') as fin, open('ariprog.out', 'w') as fout:
length = int(fin.readline().strip())
pqMax = int(fin.readline().strip())
def isValid(start, diff):
for i in range(1,length):
if start + (diff*i) not in bisquares:
return False
return True
bisquares = set()
for p in range(pqMax+1):
for q in range(pqMax+1):
bisquares.add(p**2 + q**2)
maxBS = max(bisquares) # do the calculation here so we only have to do it once
res=[]
for start in bisquares:
for diff in range(1, (maxBS-start)//(length-1)+1):
if start + (diff*(length-1)) > maxBS:
break
if isValid(start, diff):
heapq.heappush(res, (diff, start)) # diff goes first in the tuple so it sorts by smallest diff first
if res: # if res is non-empty
while res:
diff, start = heapq.heappop(res)
fout.write(f"{start} {diff}\n")
else:
fout.write("NONE\n")
ariprog.in
contains two lines, the first beingn
and the second being the largest thatp
andq
can be. Is that correct? Do you have a bound on how big those numbers can be? \$\endgroup\$