UVa 11191 Problem F Perfect Square

UVa 11191 Problem F Perfect Square
http://uva.onlinejudge.org/external/111/11191.html

問題概要

n(n<200000)個の数字(abs(n) < 10^18)が与えられる。
それらの数字はすべて素因数の大きさが30以下である。
それらの数字を2個選んだとき、その積が整数の平方となる数字の選び方を答えよ。
また、3個選んだときの数も求めよ。

解法

数字をどの素因数が奇数回でてるか偶数回でてるかで分類する。
さらに、正か負で分類する。
それをbitで表し、それぞれの分類について何個入力データで来たか記録する。
1bit目を正負、2bit目を素因数の2の偶奇、3bit目を3の....としていくと、
11bitで全て分類できる。
掛け算することは、それらのbitのxorをとることになる。
xorをとった結果が0になれば、それはなにかの数の二乗である。（素因数がすべて偶数であることなので）
2個選ぶときは、自分自身をxorとれば0なるし、
3個選ぶ時はとりあえずなにか2個選んで、それらのxorをとった結果に
分類される数字をかければ、0になる。
基本的な方針はこれでok.
0が入力にふくまれるので、0はなにを書けても0なことから、例外処理が必要。
また、bitで分類した結果、それが0であらわされるような数字は少し別処理が必要。
そこらへんは頑張る。

ソース

#include<iostream>

#define REP(i,b,n) for(int i=b;i<(int)n;i++)
#define rep(i,n)   REP(i,0,n)

using namespace std;
typedef long long lli;

const int baseP[] = {-1, 2, 3, 5, 7, 11, 13, 17, 19, 23, 29};
const int P = 11;

lli bitP(lli tmp){
  int ret = 0;
  if(tmp < 0){
    ret = 1;
    tmp *= -1;
  }
  REP(i, 1, P){
    bool b = false;
    while(tmp % baseP[i] == 0){
      b = !b;
      tmp /= baseP[i];
    }
    if(b){
      ret += (1 << i);
    }
  }
  return ret;
}

lli calcX(int n, int z){
  if(z == 0)return 0;
  if(n < 2)return 0;
  if(n - z < 2)return (lli)n*(n-1)/2;
  return ((lli)n * (n-1)/2 - (lli)(n-z) * (n-z-1)/2);
}
lli calcY(int n, int z){
  if(z == 0)return 0;
  if(n < 3)return 0;
  if(n-z<3){
    return (lli)n * (n-1) * (n-2)/6;
  }
  return ((lli)n * (n-1) * (n-2)/6 - (lli)(n-z) * (n-z-1) * (n-z-2)/6);
}

int main(){
  int T;
  cin >> T;
  rep(tc, T){
    lli n, zeroCnt = 0, data[(1<<P)];
    rep(i, (1<<P)){
      data[i] = 0;
    }
    cin >> n;
    rep(i, n){
      lli tmp;
      cin >> tmp;
      if(tmp == 0){
        zeroCnt++;
      }
      else {
        data[bitP(tmp)]++;
      }
    }
    lli X = calcX(n, zeroCnt), Y = calcY(n, zeroCnt);
    if(data[0] > 0)X += ((lli)data[0] * (data[0]-1) )/2;
    if(data[0] > 2)Y += ((lli)data[0] * (data[0]-1) * (data[0]-2) )/6;
    REP(i, 1, (1<<P)){
      if(data[i] == 0)continue;
      X += ((lli)data[i] * (data[i]-1))/2;
      Y += ((lli)data[i] * (data[i]-1))/2 * data[0];
      REP(j, i+1, (1<<P)){
        if((i^j) <= j)continue; 
        Y += (lli)data[i] * data[j] * data[i^j];
      }
    }
    cout << X << ' ' << Y << endl;
  }
  return 0;
}