代码源每日一题div1 平方计数

Daimayuan Online Judge

思路：计数题
注意到$A_i\le 10^6$,因此可以考虑用一个桶$a[i]$来存储数据
$ans={\sum }_{i=1}^{max{a}_i}{\sum }_{j=i+1}^{max{a}_i+1}a[i]\ast a[k]\ast ([k==j\ast j-i\ast i]\ast [k<m])$

#include <bits/stdc++.h>
using namespace std;
#define IOS ios::sync_with_stdio(false),cin.tie(0),cout.tie(0);
#define endl '\n';
typedef long long ll;
typedef pair<int,int> PII;
const int N=1e6+10,mod=1e9+7;
int n,a[N];
int main(){
    IOS;
    cin>>n;
    int m=0;
    for(int i=1;i<=n;i++){
        int t;cin>>t;
        a[t]++;
        m=max(m,t);
    }
    ll ans=0;
    for(int i=1;i<=m;i++){
        for(int j=i+1;j<=m+1;j++){
            ll k=1LL*j*j-1LL*i*i;
            if(k>m) break;
            ans+=1LL*a[i]*a[k];
        }
    }
    cout<<ans<<endl;
    return 0;
}