平方计数

传送门

题目描述

给 \(N\) 个正整数，第 \(i\) 个数用 \(a_i\) 来表示，求出有多少对 \((i,j)\) 使得 \({a_i}^2 + a_j\) 是一个完全平方数.
\(1 \le N \le 10^6, 1 \le a_i \le 10^6\)

CODE

#include <bits/stdc++.h>
#define rep(i, a, b) for(int i = (a); i <= (b); i ++ )

using namespace std;

typedef long long LL;
typedef pair<int, int> PII ;

template <typename T> void chkmax(T &x, T y) { x = max(x, y); }
template <typename T> void chkmin(T &x, T y) { x = min(x, y); }

/*
  a_i * a_i + a_j = x * x
  a_j = (x - a_i) * (x + a_i)
  a_j 有两个因此，并且其差值为 2 * a_i
  可以枚举所有的因子 i, 枚举其倍数 j 
  则两个因子分别为 i   j/i 然后判断是否符合条件即可
  若符合条件，则 a_i = d / 2 , a_j = j
*/

const int N = 1e6 + 10;
int cnt[N];

int main() { 
  int n; scanf("%d", &n);
  for(int i = 1, x; i <= n; i ++ ) {
    scanf("%d", &x); cnt[x] ++;
  }

  LL ans = 0;
  rep(i, 1, N - 1) {
    for(int j = i; j < N; j += i) {
      int mi = min(j / i, i), mx = max(j / i, i);
      int d = mx - mi;
      if(d % 2 == 0) {
        ans += 1LL * cnt[j] * cnt[d / 2];
      }
    }
  }
  printf("%lld\n", ans / 2);
  return 0;
}