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:
    ariLen = int(fin.readline().strip())
    ariLim = int(fin.readline().strip())
def generate(bound):
    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:
big_mama_ar.sort(key=lambda x: x[1])
with open('ariprog.out', 'w') as fout:
    if big_mama_ar == []:
        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



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

Welcome to CodeReview!


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.


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

Redundant pass

That pass isn't needed.

Use format strings


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

can be

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

Simplify some math



can be

2 * bound**2 + 1

due to operator precedence.



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

can be

if not parity[a + n*d]:

because 0 is falsey.


ariLen is more commonly written ari_len in Python.

| improve this answer | |
  • 2
    \$\begingroup\$ 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. \$\endgroup\$ – Toby Speight Aug 29 '19 at 7:35

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.