Codeforces 60 D. Savior

問題

どの要素も互いに異なる項数nの数列a[i]が与えられる。
a[i],a[j]について、ある整数bが存在して、
{a[i],a[j],b} = {x^2, y^2, z^2 | x^2+y^2=z^2 かつx,y,zは互いに素}
と書けるとき、a[i]とa[j]を同一視する。

a[i]は何個の部分に分かれているか、求めよ。

制約条件

a[i]≦10^7
n≦10^6

方針

原始ピタゴラス数を効率的に列挙して、
union-findでマージする。

原始ピタゴラス数は、
n,mを互いに素かつ一方が偶数である自然数として、
x=m^2-n^2
y=2*m*n
z=m^2+n^2
とあらわすことができる。

数列の少なくとも二つがx,y,zのどれかになっていなければならないので、
yがa[i]の最大値を超えない範囲まで原始ピタゴラス数を列挙すればよいことがわかる。

a[i]の最大値をMとすれば、
2*n*n≦y=2*m*n≦M
であるので、2*n*n≦Mなるnを列挙すれば十分。

一方でmは、a[i]+a[j]=zとなるi,jが存在するzまでを列挙したいので、
z≦2*M
すなわちm*m+n*n≦2*Mを満たすmを列挙すれば十分。

int N,a[1000000];
int p[1000000];
int idx[10000001];
int root(int x){
	if(x==p[x])return x;
	return p[x]=root(p[x]);
}
void merge(int a,int b){
	a=root(a); b=root(b);
	if(a!=b)p[a]=b;
}

void run()
{
	cin>>N;
	rep(i,N)cin>>a[i], p[i]=i, idx[a[i]]=i;
	
	for(int n=1;2*n*n<=10000000;n++)for(int m=n+1;2*m*n<=10000000&&m*m+n*n<=20000000;m+=2)if(__gcd(m,n)==1){
		int x=m*m-n*n, y=2*m*n, z=m*m+n*n;
		if(__gcd(x,__gcd(y,z))!=1)continue;
		int p=x<=10000000?idx[x]:0;
		int q=y<=10000000?idx[y]:0;
		int r=z<=10000000?idx[z]:0;
		if(a[q]==y&&a[r]==z)merge(q,r);
		if(a[p]==x&&a[r]==z)merge(p,r);
		if(a[p]==x&&a[q]==y)merge(p,q);
	}
	set<int> s;
	rep(i,N)s.insert(root(i));
	cout<<s.size()<<endl;
}