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

Question

我正在尝试解决 USACO 的算术级数问题。这是问题陈述。

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.

我已经实现了一种可以正确解决问题的算法，但花费的时间太长。在 n = 25 和 m = 250 的最大约束下，我的程序无法在 5 秒的时间限制内解决问题。

这是代码:

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))

程序输出了正确答案，但花费的时间太长。我尝试使用最大 n/m 值运行该程序，30 秒后，它仍在运行。

Answer 1

免责声明:这不是完整的答案。这更多的是寻找的一般方向。

对于序列的每个成员，您需要寻找四个参数:两个要进行平方和求和的数字( q_i 和 p_i )，以及要在下一步中使用的两个差值( x 和y ) 这样

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

主题:

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

有太多的未知数，所以我们无法得到封闭式的解决方案。

• 让我们修复 b :(还有太多未知数)
• 让我们修复 q_i ，并声明这是序列的第一个成员。即，让我们从 q_1 = 0 开始搜索，尽可能扩展，然后提取长度为 n 的所有序列。尽管如此，还有太多的未知数。
• 让我们修复 x :我们只有p_i和y来解决。此时，请注意，满足方程的可能值的范围远小于 0..m 的整个范围。 。经过一些微积分，b = x*(2*q_i + x) + y*(2*p_i + y) ，并且确实没有太多值需要检查。

最后一步修剪是它与完整搜索的区别。如果你明确地写下这个条件，你可以得到可能的范围 p_i值并从中找到可能序列的长度，步骤 b作为 q_i 的函数和x 。拒绝小于 n 的序列应进一步修剪搜索。

这应该让你从 O(m**4)复杂性〜O(m**2) 。进入时间限制应该足够了。