USACO Training - Section 1.5 - Arithmetic Progressions [closed]

Question

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:

5  
7

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: ***
LANG: PYTHON3
TASK: ariprog
'''

#import os

#try:
#    os.chdir(os.path.dirname(__file__))
#except:
#    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:
            break
        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])
else:
    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.

Answer 1

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)

main()

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

Answer 2

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!