Week4 作业 B - 四个数列【二分】

题目描述

有四个数列 A,B,C,D，每个数列都有 n 个数字。ZJM 从每个数列中各取出一个数，多少种方案使得 4 个数的和为 0？（当一个数列中有多个相同的数字的时候，把它们当做不同的数对待）。

输入

• 第一行：n（代表数列中数字的个数） （1≤n≤4000）

• 接下来的 n 行中，第 i 行有四个数字，分别表示数列 A,B,C,D 中的第 i 个数字（数字不超过 2 的 28 次方

输出

• 输出不同组合的个数

思路

• 首先看数据规模，四个数列，每个最多4000个数。最简单的穷举方法显然会超时。既然全部穷举不行，我们可以考虑部分穷举。
• 最后答案的值可能超过int的存储范围，所以最好使用long long。

假设从四个数列中取出的数是a，b，c，d。a+b+c+d=0 也即 a+b=-(c+d)。因此，我们不妨分两次穷举。首先穷举A,B，把这两个数列的和都列出来存在一个数组sum1中。然后再穷举C、D，看看这两个数列中的数的和的负数和sum1中的有多少相等。

这样做本质上是空间换时间。既然存在了sum1中，我们便可以对sum1进行排序(复杂度$O(n^{2}log(n^2))$，可以接受)，然后二分查找。

注意

对于每一个-(c+d)对应的a+b的值可能不只一个，所以，要进行两次二分查找（可以分别查找sum中第一个满足条件的值和最后一个满足条件的值对应的索引，根据其差值确定满足条件的值的数量）。

完整代码

#include<iostream>
#include<algorithm>
using namespace std;
int n;
int sum1[16000005];
int a[4005], b[4005], c[4005], d[4005];

int find(int x)
{
	int l=0, r=n*n;
	int p1, p2;
	while(r > l)
	{
		p1 = (r+l)>>1;
		if(sum1[p1] >= x)	r = p1;
		else l = p1+1;
	}
	
	int l2=0, r2=n*n;
	while(r2 > l2)
	{
		p2 = (r2+l2)>>1;
		if(sum1[p2] <= x)	l2 = p2+1;
		else r2 = p2;
	}
	
	if(sum1[l] != x)	return 0;
	else return (l2 - l);
}

int main ()
{
	cin>>n;
	for(int i=0;i<n;i++)	cin>>a[i]>>b[i]>>c[i]>>d[i];
	
	int cnt=0;
	for(int i=0;i<n;i++)
	{
		for(int j=0;j<n;j++)
		{
			sum1[cnt++] = a[i] + b[j];
		}
	}
	sort(sum1, sum1+n*n);
	
	long long res=0;
	for(int i=0;i<n;i++)
	{
		for(int j=0;j<n;j++)
		{
			res += find(-(c[i]+d[j]));
		}
	}
	
	cout<<res<<endl;
	return 0;
}