如何优化 (3*O(n**2)) + O(n) 算法？

Question

I am trying to solve the arithmetic progression problem from USACO. Here is 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.

I have implemented an algorithm which correctly solves the problem, yet it takes too long. With the max constraints of n = 25 and m = 250, my program does not solve the problem in the 5 second time limit.

Here is the code:

n = 25
m = 250

bisq = set()
for i in range(m+1):
    for j in range(i,m+1):
        bisq.add(i**2+j**2)

seq = []
for b in range(1, max(bisq)):
    for a in bisq:
        x = a
        for i in range(n):
            if x not in bisq:
                break
            x += b
        else:
            seq.append((a,b))

The program outputs the correct answer, but it takes too long. I tried running the program with the max n/m values, and after 30 seconds, it was still going.

Answer 1

Disclaimer: this is not a full answer. This is more of a general direction where to look for.

For each member of a sequence, you're looking for four parameters: two numbers to be squared and summed ( q_i and p_i ), and two differences to be used in the next step ( x and y ) such that

q_i**2 + p_i**2 + b = (q_i + x)**2 + (p_i + y)**2

Subject to:

• 0 <= q_i <= m
• 0 <= p_i <= m
• 0 <= q_i + x <= m
• 0 <= p_i + y <= m

There are too many unknowns so we can't get a closed form solution.

• let's fix b : (still too many unknowns)
• let's fix q_i , and also state that this is the first member of the sequence. Ie, let's start searching from q_1 = 0 , extend as much as possible and then extract all sequences of length n . Still, there are too many unknowns.
• let's fix x : we only have p_i and y to solve for. At this point, note that the range of possible values to satisfy the equation is much smaller than full range of 0..m . After some calculus, b = x*(2*q_i + x) + y*(2*p_i + y) , and there are really not many values to check.

This last step prune is what distinguishes it from the full search. If you write down this condition explicitly, you can get the range of possible p_i values and from that find the length of possible sequence with step b as a function of q_i and x . Rejecting sequences smaller than n should further prune the search.

This should get you from O(m**4) complexity to ~ O(m**2) . It should be enough to get into the time limit.

Answer 2

A couple more things that might help prune the search space:

b <= 2*m*m//n

a <= 2*m*m - b*n

An answer on math.stackexchange says that for a number x to be a bisquare, any prime factor of x of the form 3 + 4k (eg, 3, 7, 11, 19, ...) must have an even power. I think this means that for any n > 3, b has to be even. The first item in the sequence a is a bisquare, so it has an even number of factors of 3. If b is odd, then one of a+1b or a+2b will have an odd number of factors of 3 and therefore isn't a bisquare.