1857F - Sum and Product - CodeForces Solution

#include <iostream>
#include <algorithm>
#include <cstring>
#include <queue>
#include <set>
#include <vector>
#include <cmath>
#include <unordered_map>
#include <map>
#define x first
#define y second

using namespace std;

typedef long long LL;

typedef pair<LL,LL> PII;

typedef pair<PII,int> PIII;

typedef pair<int,string> PIS;

const int N = 200010,M = N * 2,mod = 1e9 + 7;

int n,m;
int f[N],g[N];

PII get(LL a,LL b)
{
	LL p = sqrt(abs(a * a - b * 4));
	if(p * p != abs(a * a - b * 4)) return {0,0};
	return {(p - a) / 2,(p + a) / 2};
}

int main(){
	int T;
	cin >> T;
	while(T--)
	{
		map<int,int> mp;
		scanf("%d",&n);
		for(int i = 1;i <= n;i++) scanf("%d",&f[i]),mp[f[i]]++;
		
		scanf("%d",&m);
		while(m--)
		{
			LL a,b;
			cin >> a >> b;
			PII k = get(a,b);
			LL res = 0;
			if(mp.count(-(k.x)) != 0 && mp.count(k.y) != 0)
			{
				if(-(k.x) == k.y)
				{
					int p = mp[k.y];
					res += (LL)p * (p - 1) / 2;
				}
				else res += (LL)mp[-(k.x)] * mp[k.y];
			}
//			if(mp.count(k.x) != 0 && mp.count(-(k.y)) != 0)
//			{
//				if(-(k.y) == k.x)
//				{
//					int p = mp[k.x];
//					res += (LL)p * (p - 1) / 2;
//				}
//				else res += (LL)mp[k.x] * mp[-(k.y)];
//			}
			cout << res << '\n';
		}
		cout << '\n';
	}
	return 0;
}