【代码源每日一题Div1】平方计数「nlogn巧妙枚举倍数优化」

平方计数

题目描述：

给你n个数字，求存在多少对(i, j)，满足 $a_i^2+a_j$是一个完全平方数

思路：

假设 $a_i^2+a_j=x^2$

则 $a_j=(x+a_i)\ast (x-a_j)$

即存在a_j的因子p,q，p*q=a[j]，p-q = 2*a[i]

显然一个简单的思路是去枚举a[j]的所有因数，看存在多少个满足条件的，但是复杂度是 $O(n\sqrt{n})$，会超时

我们需要考虑进行一下优化

可以发现，其实就是利用了两个因数，而且这两个因数的乘积不超过1e6

所以我们可以通过枚举每个数字i的倍数的方法去优化，即

for(int i = 1; i <= 1000000; ++i){
     
     
		for(int j = i; j <= 1000000; j += i){
     
     //j是其中的一个因数，j/i是另一个因数
			int mi = min(i, j/i), ma = max(i, j/i);
			int d = ma - mi;
			if(d%2==0){
     
     
				ans += tr[j]*tr[d/2];
			}
		}
	}

最后输出ans/2就行

#include<bits/stdc++.h>
using namespace std;
#define int long long
#define endl '\n'
#define io ios::sync_with_stdio(false), cin.tie(0), cout.tie(0)
#define m_p(a, b) make_pair(a, b)

typedef long long ll;
typedef pair<int, int> pii;

#define MAX 1000050

int n, m, x;
int tr[MAX];

void work(){
    
    
	cin >> n;
	ll ans = 0;
	for(int i = 1; i <= n; ++i){
    
    
		cin >> x;++tr[x];	
	}
	for(int i = 1; i <= 1000000; ++i){
    
    
		for(int j = i; j <= 1000000; j += i){
    
    
			int mi = min(i, j/i), ma = max(i, j/i);
			int d = ma - mi;
			if(d%2==0){
    
    
				ans += tr[j]*tr[d/2];
			}
		}
	}
	cout << ans/2 << endl;
}

signed main(){
    
    
	io;
	work();
	return 0;
}