# USACO Arithmetic Progressions

The problem I'm solving is:

given a length of an arithmetic progression, and a limit to the terms (see below), find all progressions that work.

All terms of the progressions must be of the form a²+b², and 0 ≤ a²+b² ≤ limit.

I have the following code:

from math import ceil
with open('ariprog.in') as fin:
def generate(bound):
max_len=((bound**2)*2)+1
parity = [0]*max_len
for i in range(bound+1):
for j in range(bound+1):
parity[i**2+j**2] = 1
return parity

parity = generate(ariLim)
lenpar = len(parity)
big_mama_ar = []
# print(lenpar)
for a in range(lenpar-1):
if parity[a] == 1:
for d in range(1, ceil((lenpar-a)/(ariLen-1))):
for n in range(1, ariLen):
# print('a:', a)
# print('d:', d)
# print('n:', n)
if parity[a+n*d] != 1:
break
else:
big_mama_ar.append((a,d))
pass
big_mama_ar.sort(key=lambda x: x[1])
with open('ariprog.out', 'w') as fout:
if big_mama_ar == []:
fout.write('NONE\n')
else:
for i in big_mama_ar:
fout.write(str(i[0])+' '+str(i[1])+'\n')


This code times out on my grader when ariLen is 21 and ariLim is 200. The time limit is 5 seconds, and on my computer, it takes 22 seconds. ariprog.in is

21

200

• How can we run this code? What's in that ariprog.in file? – Georgy Aug 28 '19 at 18:29

Welcome to CodeReview!

## Whitespace

The PEP8 standard, and consequently most Python linting tools, will recommend that you add another linebreak before your function definitions, plus some whitespace around your operators, etc. I won't detail this exhaustively; you're best to use the IDE of your choice - PyCharm is a reasonable one that is helpful for this.

## Type hinting

bound is probably an integer, so add : int after it. It probably returns an int as well.

## Subroutines

Put your global-scoped code into subroutines for ease of maintenance, legibility and testing.

## Redundant pass

That pass isn't needed.

## Use format strings

This:

str(i[0])+' '+str(i[1])+'\n'


can be

f'{i[0]} {i[1]}\n'


## Simplify some math

This:

((bound**2)*2)+1


can be

2 * bound**2 + 1


due to operator precedence.

## Truthiness

This:

if parity[a+n*d] != 1:


can be

if not parity[a + n*d]:


because 0 is falsey.

## camel_case

ariLen is more commonly written ari_len in Python.

• Not just "because 0 is falsey" - but because 1 is the only truthy value we'll get. If parity could be 2, for example, then those tests wouldn't be equivalent. – Toby Speight Aug 29 '19 at 7:35