Wannafly挑战赛27 A-灰魔法师(思维)

#include <bits/stdc++.h>
#define maxn 100000
#define ll long long
using namespace std;
int pre[maxn],a[maxn];
int n,num;
  
int main()
{
  scanf("%d",&n);
  num = 1;    // 不知道为啥不能从0开始...
  for(int i=1;i*i<=200000;i++){
    a[num++] = i * i;
  }
  memset(pre,0,sizeof(pre));
  for(int i=0;i<n;i++){
    int x;
    scanf("%d",&x);
    pre[x] ++;
  }
  ll sum = 0,sum1 = 0;
  for(int i=1;i<=num;i++){
    for(int j=1;j<=maxn;j++){
      if(pre[j] == 0) continue;
      int ans = a[i] - j;
      if(ans > maxn) continue;
      if(ans <= 0) break;
      if(pre[ans]){
        if(ans == j) sum1 += ((ll)pre[j] * (ll)(pre[j] - 1)) / 2;
        else sum += (ll)pre[ans] * (ll)pre[j];    // 这里不要除以2，会有精度损失
      }
    }
  }
  printf("%lld\n",sum1 + sum / 2);
  return 0;
}