USACO Section 1.4 Arithmetic Progressions

/*
ID: lucien23
PROG: ariprog
LANG: C++
*/

#include<iostream>
#include<fstream>
#include<set>
#include<vector>
#include<algorithm>
using namespace std;
typedef struct{
int first;
int different;
}ArithProgressions;

bool compProgressions(ArithProgressions item1,ArithProgressions item2)
{
if(item1.different<item2.different)
return true;
else if(item1.different>item2.different)
return false;
else
{
if(item1.first<item2.first)
return true;
else
return false;
}
}

int main()
{
ifstream infile("ariprog.in");
ofstream outfile("ariprog.out");
if(!infile || !outfile)
{
cout<<"file operation failure!"<<endl;
return -1;
}
int N,M;
infile>>N>>M;

set<int> nums;
vector<ArithProgressions> progressions;

for (int i=0;i<=M;i++)
{
for (int j=0;j<=M;j++)
{
nums.insert(i*i+j*j);
}
}

for (set<int>::iterator it=nums.begin();it!=nums.end();it++)
{
set<int>::iterator it1=it;
while(++it1!=nums.end())
{
int diff=*it1-*it;
if(*it+(N-1)*diff>2*M*M)
break;
int i;
for (i=2;i<N;i++)//此处i不能从0开始，否则会超时
{
int item=*it+i*diff;
if(nums.find(item)==nums.end())
{
break;
}
}
if(i==N)
{
ArithProgressions progression;
progression.first=*it;
progression.different=diff;
progressions.push_back(progression);
}
}
}

sort(progressions.begin(),progressions.end(),compProgressions);

if(progressions.size()==0)
outfile<<"NONE"<<endl;
else
{
for (vector<ArithProgressions>::iterator it=progressions.begin();it!=progressions.end();it++)
{
outfile<<it->first<<" "<<it->different<<endl;
}
}

return 0;
}