AtCoder Beginner Contest 342, D : Square Pair

問題文

　https://atcoder.jp/contests/abc342/tasks/abc342_d

問題概要

　$n$ 項からなる非負自然数の列 $\langle A_1, A_2, \dots, A_n \rangle$ が与えられる．以下の条件を満たす順序対 $( i, j )$ はいくつあるか？

• $1 \leq i < j \leq n$
• $A_i A_j$ は平方数

制約

• $2 \leq n \leq 2 \times 10^5$
• $0 \leq A_i \leq 2 \times 10^5$

解法

　平方数の性質としてまず確認しておきたいことがあります．それは，

• ある数が平方数 $\Leftrightarrow$ その数を素因数分解したとき，各素因数の肩に現れる指数が全て偶数

というものです．$0$ でない 2 数をかけると指数は和になるので，指数が奇数であるような素因数の集合が等しい数同士をかけると，積が平方数になります．$0$ をかけるとなんでも平方数になってしまいますが*1，特殊ケースということにして後で考えます．正整数同士の積については前述のように指数が奇数な素因数の集合が分かればよいので，それぞれ素因数分解して各素因数の個数を数えるなどして変換し，個数を数えておけば計算できます．具体的には，ある（奇数乗の）素因数集合が $k$ 個あるとき，そこから 2 つ選ぶ方法は $$\frac{ k ( k - 1 ) }{ 2 }$$ 通りとなります．
　続いて少なくとも片方が $0$ になる組合せの数え方ですが，「少なくとも片方が $0$」は「両方が正」の余事象なので，後者の個数を全体から引くことで求められます．なので，$A$ 中の正整数が $m $ 個のとき，積が $0$ になる選び方は $$\frac{ n ( n - 1 ) }{ 2 } - \frac{ m ( m - 1 ) }{ 2 }$$ 通りとなります．
　あとはこれらの値を集計すれば，答えが求まります．
　計算量としては，まず $n$ 個の数を素因数分解するのに $O( n \sqrt{ \max A } )$ 時間かかります．素因数の集合を std::map で管理することにすれば，要素数が $O( n )$ 個でキーの比較に $O( \log \max A )$ 時間かかる*2ので，$O( n \log \max A )$ 時間がかります．あとは正整数を数えるのに $\Theta( n )$ 時間かかったりしますが定数倍の中に消えるので，全体では $O( n ( \sqrt{ \max A } + \log \max A ) )$ 時間となります．

コード

// 素因数分解 O( √N )
vector< long long > primeFactorization( long long N );
// 中身省略

int main()
{
	IN( int, N );
	VI A( N );
	cin >> A;

	VVI vecs( N );
	map< VI, int > counts;
	REP( i, N )
	{
		const auto factors = primeFactorization( A[i] );
		map< int, int > cnt;
		for_each( ALL( factors ), [&]( const int f ){ ++cnt[f]; } );

		VI vec;
		FOR( p, cnt )
		{
			if ( p.snd % 2 )
			{
				vec.PB( p.fst );
			}
		}

		vecs[i] = vec;
		++counts[ vec ];
	}

	LL res = 0;
	REP( i, N )
	{
		if ( A[i] )
		{
			res += counts[ vecs[i] ] - 1;
		}
	}
	res /= 2;

	const LL nonzeros = count_if( ALL( A ), bind( greater< int >(), _1, 0 ) );
	res += ( LL( N ) * ( N - 1 ) - nonzeros * ( nonzeros - 1 ) ) / 2;

	cout << res << endl;

	return 0;
}

*1:$0 = 0^2$ なので $0$ は平方数．

*2:素因数の個数の最大値を $O( \log \max A )$ で抑える気持ち．