I am working on the USACO problem Arithmetic Progression. Here is their 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 p² + q² (where p and q are non-negative integers).

Furthermore, input is read from a file ariprog.in, and output should be written to another file ariprog.out. Help is not required for this part, but a sample ariprog.in is shown below:


And here is the input format:

Line 1: N (3≤N≤25), the length of progressions for which to search
Line 2: M (1≤M≤250), an upper bound to limit the search to the bisquares with 0≤p,q≤M.

And here is the output format:

If no sequence is found, a single line reading `NONE'. Otherwise, output one or more lines, each with two integers: the first element in a found sequence and the difference between consecutive elements in the same sequence. The lines should be ordered with smallest-difference sequences first and smallest starting number within those sequences first.

There will be no more than 10,000 sequences.

For my code, here are my steps.

  1. Read the file and save variables N and M
  2. Generate all possible bi-squares (p² + q²) and save it to a list
  3. Loop through starting values for arithmetic progressions
  4. Loop through ending values for arithmetic progressions
  5. If the end of a possible sequence (calculated with start + (length - 1) * difference is larger than the largest value in the list of possible bi-squares (found in step 2), skip this case
  6. Test a starting value and difference by looping through it until the length is achieved
  7. Sort the list according to output format
  8. Write to file

Out of these, steps 3-6 take the time. I am interested in speed improvements and optimizations, not a full solution.

My logic is working, I pass 5/9 of the tests. I fail on the test case with the input being 18 100 due to too much time taken

Now, here is the actual code so far:

ID: ***
TASK: ariprog

#import os

#    os.chdir(os.path.dirname(__file__))
#    pass

#import time

from itertools import combinations_with_replacement

with open('ariprog.in', 'r') as file:
    N, M = map(int, [line.replace('\n', '') for line in file.readlines()])
def calculateBisquares(limit):
    p_or_q_values = range(limit+1)
    return sorted(list(set([((i[0]**2) + (i[1]**2)) for i in combinations_with_replacement(p_or_q_values, 2)])))

def checkNum(bisquare_index, all_bisquares, difference):
    for sequence_item_number in range(2, N):
        if(int(all_bisquares[bisquare_index] + difference * sequence_item_number) not in all_bisquares):
            return False
    return True

#print(f'starting: {time.time()}')
#start = time.time()

all_bisquares = calculateBisquares(M)
maximum_number = 2 * (M ** 2)
sequences = []

for starting_index in range(0, len(all_bisquares) - N + 1):
    for ending_index in range(starting_index + 1, len(all_bisquares) - N + 2):
        difference = all_bisquares[ending_index] - all_bisquares[starting_index]

        if all_bisquares[starting_index] + (N - 1) * difference > maximum_number:
        if checkNum(starting_index, all_bisquares, difference):
            sequences.append([all_bisquares[starting_index], difference])

if len(sequences) != 0:
    sequences.sort(key=lambda d: d[1])

    out = '\n'.join([' '.join(list(map(str, sequence))) for sequence in sequences])
    out = 'NONE'

#print(f'ending: {time.time()}')
#print(f'time taken: {time.time()-start}')

with open('ariprog.out', 'w') as file:
    file.write(out + '\n')

Side note: I am interested in an optimization but haven't yet implemented in Python.

  • \$\begingroup\$ Sounds like you already know what to improve (given your last sentence). I recommend you implement that optimisation and return with your improved code. \$\endgroup\$ Jul 21 at 6:45
  • \$\begingroup\$ Although I know what to improve, I am not entirely sure how this can be implemented in Python with my above algorithm. Can you provide some guidance? \$\endgroup\$
    – cold10
    Jul 21 at 14:36
  • 1
    \$\begingroup\$ This is Code Review - asking for help with implementation is specifically off-topic, I'm afraid. \$\endgroup\$ Jul 22 at 14:13

2 Answers 2


The link you provided for an optimization suggested using a Boyer-Moore like algorithm. That got me thinking about using string.find() and then re.search() to find the sequences. The idea is to create a bytearray in which the bisquares are marked with b'1' and the other numbers are marked with b'0'. Then build a regular expression to search the bytearray for a sequence of equally-spaced b'1's of the desired length.

import re

def bisquares(n):
    limit = 2*n*n
    bsqs = bytearray(b'0' * (limit + 1))
    for i in range(n+1):
        for j in range(i, n+1):
            k = i*i+j*j
            if k <= limit:
                bsqs[k] = ord('1')
    return bsqs

def find_arithmetic_sequences(n, m):
    bss = bisquares(m)

    for width in range(len(bss)//n + 1):
        # pattern looks like b'1..1..1..1, where 'width' determines
        # the number of '.'s and 'n' determines the number of '1's.
        pattern = f'1(?:.{{{width}}}1){{{n-1}}}'.encode('utf8')
        regex = re.compile(bytes(pattern))

        start = 0
        while (match := regex.search(bss, start)):
            start = match.start() + 1
            yield match.start(), width + 1

def main():
    with open('ariprog.in', 'r') as ifile:
        N = int(next(ifile))
        M = int(next(ifile))
    found_some = False
    with open('ariprog.out', 'w') as ofile:
        for index, difference in find_arithmetic_sequences(N, M):
            print(index, difference, file=ofile)
            found_some = True
        if not found_some:
            print('NONE', file=ofile)


For N, M = 18, 100 it takes under 2 seconds, compared to over 36 for your code.


PEP-8: consider renaming to calculate_bisquares() and check_num().

Also, definitely make it lint clean with flake8 before requesting review.

nit: In

    N, M = map(int, [line.replace('\n', '') for line in file.readlines()])

there's no need to allocate storage for a list with [ ] -- a generator expression would suffice. This only reads two lines so the effect is trivial. But it doesn't hurt to get in the habit of avoiding an extra copy which is immediately discarded, since sometimes we see big cache-busting lists.

Same thing here:

    return sorted(list(set([((i[0]**2) + (i[1]**2)) for i in combinations_with_replacement(p_or_q_values, 2)])))

The list() call is redundant, as sorted() will produce a list.

As a separate item, for i in ... forces the use of inconvenient [0] & [1] subscripts. Prefer tuple unpack:

for p, q in ...

Delete commented code.

Introduce a verbose or timing flag for conditional output of the elapsed time.

Rather than check_num(), choose an informative identifier. Begin by writing a one-sentence """docstring""" for the function. With that in hand, you'll be in a better position to summarize what the function does. Saying that it "checks" is too vague, here.

Also, perhaps the int() call is redundant?

Also, write a unit test for this function.

Starting with the all_bisquares assignment you're needlessly putting lots of variables into the global environment. Prefer to wrap def main(): around those statements. That way, local variables will go out of scope when the function exits.

Perhaps you could do Extract Helper on the for loop? Consider naming it def _get_sequences():.

I am reading this expression:

'\n'.join([' '.join(list(map(str, ...

Same remark as above. No need for a [ ] list comprehension when a generator expression would do. And no need for redundant list() call.

Rather than a pair of .join's, perhaps you would find pp pretty printing more convenient?

from pprint import pp

The structure is reasonably good, things are broken into clear bite-size pieces.

You mentioned a desire to refactor to a more efficient algorithm. Write some unit tests for these nice simple functions, and then verify your refactor speeds things up without breaking unit tests.

Consider using "import hypothesis", which is an adventure in itself!

  • \$\begingroup\$ What is Extract Helper? I did a quick Google Search and couldn't find anything relevant. \$\endgroup\$
    – cold10
    Jul 21 at 14:15
  • 1
    \$\begingroup\$ Sure, that's fine. Let's start with another defined term, "refactor". It means to change a program so it does "the same thing", but "better", usually so it is more easily read, understood, and changed, also to improve things like performance characteristics. And then "extract helper" means to pull out a few lines of code and package them up into their own function. Then we have a nice clear cut API around it, can document it, write unit tests for it and so on. Many IDEs (interactive development environments) offer that and related refactors when you explore their menu options. \$\endgroup\$
    – J_H
    Jul 21 at 23:45

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