# Help on USACO training 1.2 Arithmetic Progressions (ariprog) TLE

Arithmetic Progressions

Unfortunately, the USACO training pages require one to be signed in, so here is a description of the problem.

Description of Arithmetic Progressions

Arithmetic Progressions

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

### 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.|

5
7

### 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.

### SAMPLE OUTPUT (file ariprog.out)

1 4
37 4
2 8
29 8
1 12
5 12
13 12
17 12
5 20
2 24

I’m using python, and whatever I do, I can only get past 7 of the 9 test cases.

My code finds all bisquares, then iterates through them and checks for possible arithmetic progressions. Finally, it outputs the arithmetic progessions.

Here is my code:

``````"""
ID: shawn.z2
LANG: PYTHON3
"""
import sys

sys.stdin = open('ariprog.in', "r")
sys.stdout = open('ariprog.out', "w")

m2 = 2 * M * M
found_ariprogs = []
bisquares = set()
lookup_table = [False] * (m2 + 1)

# find bisquares
for i in range(M + 1):
i2 = i ** 2
for j in range(i, M + 1):
bisquare = i2 + j ** 2
lookup_table[bisquare] = True

bisquares = sorted(list(bisquares))

def is_ariprog(a, b):
# checks each term in arithmetic progression
for n in range(N):
if not lookup_table[a + n * b]:
return False
return True

# iterate through bisquares and check for arithmetic progressions
l_bisquares = len(bisquares)
for i in range(l_bisquares):
a = bisquares[i]
lim = (m2 - a) // (N - 1)
# for each bisquare where the difference between terms in the possible arithmetic progression doesn't make the last term in the arithmetic progression go over the last bisquare
# test the arithmetic progression
for j in range(i + 1, l_bisquares):
b = bisquares[j] - a
if b > lim:
break
if is_ariprog(a, b):
found_ariprogs.append((b, a))

# output
found_ariprogs.sort()
if len(found_ariprogs) == 0:
sys.stdout.write("NONE\n")
sys.exit(0)
for ariprog in found_ariprogs:
sys.stdout.write(f"{ariprog} {ariprog}\n")
``````

The specific test case I’m stuck on is the 8th test case, where the input is:
22
250
and output is:
13421 2772

I always time out on this test case.
Can anyone help? Thanks!

I think it just might be because Python is too slow.

1 Like time to switch to java ig