代码源：607、平方计数本文已参与[新人创作礼]活动，一起开启掘金创作之路 Offer 驾到，掘友接招！我正在参与202

本文已参与[新人创作礼]活动，一起开启掘金创作之路

Offer 驾到，掘友接招！我正在参与2022春招打卡活动，点击查看活动详情。

题目描述

这是3月21日代码源div1的每日一题。

平方计数 - 题目 - Daimayuan Online Judge

题目描述

给N个正整数, 第ii个数用ai来表示, 求出有多少对(i,j)(i,j) 使得ai^2+aj是一个完全平方数.

输入格式

第一行一个正整数n 第二行n个数, 表示a1,a2,a3...an.

输出格式

一行一个整数, 表示答案

数据范围

对于所有测试数据满足1≤n≤10^6, 且1≤ai≤10^6

样例输入

5
1 2 3 4 5

样例输出

看到这题第一反应，枚举ai，再枚举bi，计算ai^2+bi是不是完全平方数，但很显然1e6的数据要超时。遂想到能不能反向逆着来，我枚举可能的完全平方数，然后根据ai去求bi，但一开始想法有问题，觉着复杂度和枚举bi差不多，幸得群里大佬（某墨镜白猫猫）解惑才想明白。

我们先预处理，数组f记录下每个数的出现次数，顺便记录下最大的max_num（作为枚举时的上界），然后从1开始枚举a，一直枚举到max_num，在来枚举完全平方数。因为a^2+b=c^2，那么b=c^2-a^2。所以我们从a开始枚举c，之前记录的最大值max_num就是b的上界（最大也不会大过max_num）。在此区间内枚举c，求出b后，再根据之前记录的a的数量和b的数量，看能有多少数对组合（其实就是f[a]*f[b]个）。

然后我们说下为什么枚举完全平方数会不超时，比如a是1开始，b最大取到1e6，

即c^2-a^2=1e6 ——> c^2-1<=1e6 ——> c=1e3，由此可以看出，即使是最大的情况，c也仅仅枚举了1000次这样，而且随着a的增长，c枚举的范围也会变小。

如果a从1枚举到1e6,b最大取到1e6，那么c总的枚举次数就是：

c^2-a^2<=1e6

a<=1e6;

a^2<=1e12;

c^2-1e12<=1e6;

c^2<=1e12+1e6;

c 约= 4*1e6

这样一算，我们总程序的只用跑4*1e6就可以得出结果，而c++1秒内可以跑的数据大概是1e8到1e9这样，显然不会超时。

AC代码

#include<iostream>
using namespace std;
#include<vector>
#include<algorithm>
#include<math.h>
#include<set>
#include<numeric>
#include<string>
#include<string.h>
#include<map>
#include<unordered_map>
#include<stack>
#include<list>
#include<queue>
#include<iomanip>


#pragma GCC optimize(1)
#pragma GCC optimize(2)
#pragma GCC optimize(3,"Ofast","inline")

#define endl '\n';
typedef long long ll;
typedef pair<int, int>PII;
const int N = 1e6 + 50;
int f[N];

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    int n;
    cin >> n ;
    vector<int>v(n);
    int max_num = 0;
    for (int i = 0; i < n; i++)
    {
        cin >> v[i];
        max_num = max(max_num, v[i]);
        f[v[i]]++;
    }
    int res = 0;
    for (int i = 1; i <= max_num; i++)
    {
        if (f[i] == 0)continue;
        for (int j = i + 1; (j - i) * (j + i) <= max_num;j++)
        {
            res += f[i]*f[(j - i) * (j + i)];
        }
    }
    cout << res << endl;
    return 0;
}