The 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 p2 + q2 (where p and q are non-negative integers).

The two lines of input are n and m, which are the length of each sequence, and the upper bound to limit the search of the bi squares respectively.

My code solves the problem correctly, however it is too slow for the 5 second time constraint. Is there any way to further optimize my code?

import heapq

with open('ariprog.in', 'r') as fin, open('ariprog.out', 'w') as fout:
    length = int(fin.readline().strip())
    pqMax = int(fin.readline().strip())

    def isValid(start, diff):
        for i in range(1,length):
            if start + (diff*i) not in bisquares:
                return False
        return True

    bisquares = set()
    for p in range(pqMax+1):
        for q in range(pqMax+1):
            bisquares.add(p**2 + q**2)

    maxBS = max(bisquares) # do the calculation here so we only have to do it once

    for start in bisquares:
        for diff in range(1, (maxBS-start)//(length-1)+1):
            if start + (diff*(length-1)) > maxBS:
            if isValid(start, diff):
                heapq.heappush(res, (diff, start)) # diff goes first in the tuple so it sorts by smallest diff first
    if res: # if res is non-empty
        while res:
            diff, start = heapq.heappop(res)
            fout.write(f"{start} {diff}\n")
  • \$\begingroup\$ It looks like ariprog.in contains two lines, the first being n and the second being the largest that p and q can be. Is that correct? Do you have a bound on how big those numbers can be? \$\endgroup\$
    – Teepeemm
    Jul 25, 2023 at 16:45
  • \$\begingroup\$ @Teepeemm Here are the restrictions: N (3 <= N <= 25), the length of progressions for which to search. M (1 <= M <= 250), an upper bound to limit the search to the bisquares with 0 <= p,q <= M \$\endgroup\$ Jul 25, 2023 at 19:55
  • \$\begingroup\$ Can you give me a sample input? \$\endgroup\$
    – user275173
    Jul 29, 2023 at 19:05

1 Answer 1


Some tiny things:

  • Instead of using power p**2 + q**2, use multiplication p*p + q*q.
  • When generating bisquares, use q in range(p, pqMax+1) to avoid duplicating work such as 1^2 + 2^2 = 2^2 + 1^2.

Most importantly, you can achieve a visible speedup with modulo arithmetic. I was inspired to do this after reading some output and because the sample code doesn't really use any properties of bisquares. I use brute force to show that arithmetic progressions reaching minimum lengths (4, 6, 14) must have differences that are multiples of interesting numbers (4, 12, 84).

import math

def evaluate_residues(mod_num, min_len):
    unique_residues = set()
    for p in range(mod_num):
        for q in range(p, mod_num):
            start = p * p + q * q
            unique_residues.add(start % mod_num)

    common_diff = 0
    for start in unique_residues:
        for diff in range(1, mod_num):
            valid = True
            for i in range(min_len):
                if ((start + i * diff) % mod_num) not in unique_residues:
                    valid = False

            if valid:
                common_diff = math.gcd(common_diff, diff)


evaluate_residues(8, 4)  # yields 4
evaluate_residues(8 * 9, 6)  # yields 12
evaluate_residues(8 * 9 * 49, 14)  # yields 84

These results can be summarized as the step size of the diff for loop:

step = 1
if length >= 4:
    step = 8
if length >= 6:
    step = 12
if length >= 14:
    step = 84
# ...
    for diff in range(step, (maxBS - start) // (length - 1) + 1, step):

When all of these things are put together, it reduces the time from >58 seconds to <1 second when length is 16 and pqMax is 250. When length is 14 and pqMax is 250, it went from >66 seconds to <6 seconds. It may be necessary to search through the bisquares more efficiently. Overall, I think your solution is well-written and I couldn't think of any ways to improve on it in terms of time complexity.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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