Wannafly挑战赛27 A 灰魔法师

链接：https://www.nowcoder.com/acm/contest/215/A
来源：牛客网
 

题目描述

“White shores, and beyond. A far green country under a swift sunrise.”--灰魔法师

给出长度为n的序列a, 求有多少对数对 (i, j) (1 <= i < j <= n) 满足 ai + aj 为完全平方数。

输入描述:

第一行一个整数 n (1 <= n <= 105)
第二行 n 个整数 ai (1 <= ai <= 105)

输出描述:

输出一个整数，表示满足上述条件的数对个数。

 

示例1

输入

复制

3
1 3 6

输出

复制

2

思路：首先能想到ai+aj最大和也就是2e5，同时2e5以内的平方数只有400多个，可以直接打表存下来，记为p[i]。

然后我们把每个a[i]做标记（同时统计个数），用p[i]去扫描a[i]，已知p[1]，令p[1]-ai，所得差即为另一个数aj，如果能在a[i]里找到（即a[j]被标记），说明这个平方和p[1]成立。复杂度O（1e2*1e5）。

接下来需要考虑两种情况：

1。完全相同的数组合：2 2 2    此时方案是3，我们标记数组应该是vis[2]=3，那么都是相同的数字，任选两个的方案数：C[3,2]，即（n*（n-1））/2。

2.不同的数字组合：1 1 3 3   此时方案是4，标记数组vis[1]=2，vis[3]=2，直接vis[1]*vis[3]就是方案数，由于重复计算（vis[3]*vis[1]），这个情况的答案需要/2，1情况不需要。

总之就是结果统计需要考虑相同数字的情况，情况就这两种- -。

#include <bits/stdc++.h>
#define ll long long
using namespace std;
ll vis[100010],n,a[1010],pos,ans1,ans2;
ll ans;
int main()
{
    pos=1;
    for(ll i=1;i*i<=200000;i++)a[pos++]=i*i;   //生成平方数
    while(~scanf("%lld",&n))
    {
        memset(vis,0,sizeof(vis));
        ll x;
        for(ll i=1;i<=n;i++)
        {
            scanf("%lld",&x);
            vis[x]++;   //统计个数
        }
        ans1=ans2=0;
        ll tmp;
        for(ll i=1;i<pos;i++)   //平方数数组
        {
            for(ll j=1;j<=100000;j++)   //1e5以内所有数
            {
                if(vis[j]==0)continue;  
                tmp=a[i]-j;
                if(tmp>100000)continue;   //不合法范围
                if(tmp<=0)break;   //不合法
                if(vis[tmp])
                {
                    if(tmp==j)ans1+=(vis[j]*(vis[j]-1))/2;   //相同时用组合公式
                    else ans2+=vis[tmp]*vis[j];   //不同时直接乘
                }
            }
        }
        ans1=ans1+ans2/2;
        printf("%lld\n",ans1);
    }
    return 0;
}