The BOSS Can Count Pairs

题面翻译

多组数据。

每组数据给你一个 $n$ 和两个序列 $a,b$。

求有多少个数对 $(i,j)$ 满足 $1 \le i < j \le n$ 且 $a_i \times a_j = b_i + b_j$

题目描述

You are given two arrays $ a $ and $ b $ , both of length $ n $ .

Your task is to count the number of pairs of integers $ (i,j) $ such that $ 1 \leq i < j \leq n $ and $ a_i \cdot a_j = b_i+b_j $ .

输入格式

Each test contains multiple test cases. The first line of input contains a single integer $ t $ ( $ 1 \le t \le 10^4 $ ) — the number of test cases. The description of test cases follows.

The first line of each test case contains a single integer $ n $ ( $ 2 \le n \le 2 \cdot 10^5 $ ) — the length of the arrays.

The second line of each test case contains $ n $ integers $ a_1,a_2,\ldots,a_n $ ( $ 1 \le a_i \le n $ ) — the elements of array $ a $ .

The third line of each test case contains $ n $ integers $ b_1,b_2,\ldots,b_n $ ( $ 1 \le b_i \le n $ ) — the elements of array $ b $ .

It is guaranteed that the sum of $ n $ across all test cases does not exceed $ 2 \cdot 10^5 $ .

输出格式

For each test case, output the number of good pairs.

样例 #1

样例输入 #1

3
3
2 3 2
3 3 1
8
4 2 8 2 1 2 7 5
3 5 8 8 1 1 6 5
8
4 4 8 8 8 8 8 8
8 8 8 8 8 8 8 8

样例输出 #1

2
7
1

提示

In the first sample, there are $ 2 $ good pairs:

$ (1,2) $ ,
$ (1,3) $ .

In the second sample, there are $ 7 $ good pairs:

$ (1,2) $ ,
$ (1,5) $ ,
$ (2,8) $ ,
$ (3,4) $ ,
$ (4,7) $ ,
$ (5,6) $ ,
$ (5,7) $ .

分析

$a[i]\times a[j] = b[i] + b[j]$ 可以得出 $a[i] <= \sqrt{2n}$，我们可以遍历数组a，用 f[a[i]][b[i]] 记录这一对ab的数量
首先我们考虑 $a[i] = a[j]$ 的情况:
这时 a[i] 和 b[i] 都已知，b[j] 也可以确定，记得处理这种情况的时候要记得去重
剩下的就是$a[i]\ne a[j]$的情况，假设$a[j] < a[i]$:
我们枚举 a[i]，对于每一个 a[i] 我们从小到大枚举j，符合条件就加到答案中去，反正根据桶来说，出现过的 a[j] 和 b[j] 才能有值，所以可以用枚举j来代替枚举 a[j]

代码

#include <iostream>
#include <cmath>
#include <algorithm>
using namespace std;
typedef long long ll;
const int N = 2e5 + 10;
int f[635][N];
int a[N], b[N];
int main()
{
	int t;
	scanf("%d", &t);
	while (t--)
	{
		int n;
		scanf("%d", &n);
		int lim = sqrt(2 * n);
		for (int i = 0; i < n; i++)
			scanf("%d", &a[i]);
		for (int i = 0; i < n; i++)
		{
			scanf("%d", &b[i]);
			//如果a[i]合法 桶里就增加
			if (a[i] <= lim)
				f[a[i]][b[i]]++;
		}

		//先考虑a[i] == a[j]的情况
		ll ans = 0;
		for (int i = 0; i < n; i++)
		{
			if (a[i] <= lim && a[i] * a[i] - b[i] >= 1 && a[i] * a[i] - b[i] <= n)
			{
				ans += f[a[i]][a[i] * a[i] - b[i]];
			}
		}
		//可能会找到自身，也即是构成了(i, i)这种对，要去掉
		for (int i = 2; i <= lim; i += 2)
		{
			//要保证i是偶数
			ans -= f[i][i * i / 2];
		}
		//又因为(ai, bi)会找到(aj, bj) 反之亦然 被计算了两次
		ans /= 2;

		//剩下的就是a[i] != a[j]的情况了 我们假设a[i] > a[j] 也只能假设a[i] > a[j] 不然有bug
		for (int i = 0; i < n; i++)
		{
			for (int j = 1; j <= lim && a[i] * j <= 2 * n && a[i] > j; j++)
			{
				if(a[i] * j - b[i] > 0 && a[i] * j - b[i] <= n)
				ans += f[j][a[i] * j - b[i]];
			}
		}
		printf("%lld\n", ans);
		for (int i = 0; i < n; i++)
			if (a[i] <= lim)
				f[a[i]][b[i]] = 0;
	}
}

在$ a[i] = a[j] $ 时，为什么结果要$-f[i][i * i / 2]$并且${\div} 2$？
去重，在找的过程中可能找到自己形成(i, i)这种不合法的对，有几个相同的块ab就会加几遍自己，所以要减去。而${\div} 2$是因为 b[i] 如果配对 b[j]，那么 b[j] 也会配对 b[i], 同一个对(i, j)会被加入答案两次，所以要${\div} 2$