USACO Arithmetic Progressions

Question

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 a2+b2, and 0 ≤ a2+b2 ≤ limit.

I have the following code:

from math import ceil
with open('ariprog.in') as fin:
 ariLen = int(fin.readline().strip())
 ariLim = int(fin.readline().strip())
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

• Georgy: How can we run this code? What's in that ariprog.in file?

Answer 1

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.

• Toby Speight: 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.