【USACO 1.4】Arithmetic Progressions

/*
TASK: ariprog
LANG:C++
URL:http://train.usaco.org/usacoprob2?a=PA9lOcZrdWq&S=ariprog
SOLVE：平方和最大为m*m*2，因此bq数组标记数i是否为平方和数，num数组存第i个平方和数
枚举公差q，从1到num[tol]/(n-1)，枚举起点p，从num[1]到num[tol-1]
判断数p+i*q是否为平方和数
*/
#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#define ll long long
using namespace std;
bool bq[200000];
int n,m,tol,found;
int p,q;
int num[200000];
int main() {
    freopen("ariprog.in","r",stdin);
    freopen("ariprog.out","w",stdout);
    scanf("%d%d",&n,&m);
    int cnt=0;
    for(int i=0;i<=m;i++)
        for(int j=i;j<=m;j++)
            bq[i*i+j*j]=1;
    for(int i=0;i<=m*m*2;i++)
        if(bq[i])num[++tol]=i;
    //printf("%d %d
",tol,num[tol]);
    int up=num[tol]/(n-1);
//    printf("%d
",up);
    for(q=1;q<=up;q++){
        for(int i=1;i<tol;i++){
            p=num[i];
            int ok=1;
            for(int j=1;j<n&&ok;j++)
                if(bq[p+j*q]==0)ok=0;
            if(ok){
                printf("%d %d
",p,q);
                found=1;
            }
        }
    }
    if(!found)puts("NONE");
    return 0;
}