PKU 2785-4 Values whose Sum is 0

問題概要

n個の数値が含まれる集合A,B,C,Dがある.
各集合から一つずつ数値を取り出したとき,その4つの数値の和が0になるパターンの個数を求めよ.

解法

蟻本で言うくじびきの問題.
A,Bの和とC,Dの和を全通り求めてから,二分探索.
計算量はO(n^2 log n).
同じ数値が複数含まれている場合に注意.

実装(C++)

#include <algorithm>
#include <vector>
#include <iostream>
#include <cstdio>
using namespace std;

#define REP(i,x) for(int i=0;i<(int)(x);i++)

#define bufsize 512
char buf[bufsize];
int bufcount=bufsize;
int getcharc(){
	if(bufcount==bufsize){
		fread(buf,sizeof(char),bufsize,stdin);bufcount=0;
	}
	return buf[bufcount++];
}
int readint(){
	int res=0;
	int h=1;
	char buf[1];
	*buf=0;
	*buf=getcharc();
	while(!((*buf>='0'&&*buf<='9')||*buf=='-'))*buf=getcharc();
	if(*buf=='-'){h=-1;*buf=getcharc();}
	while(*buf>='0'&&*buf<='9'){
		res*=10;res+=*buf-'0';
		*buf=getcharc();
	}
	return res*h;
}
int n;
int A[5000],B[5000],C[5000],D[5000];
int AB[5000*5000],CD[5000*5000];
int main(){
	setvbuf(stdin,NULL,_IOFBF,256);
	n=readint();
	REP(i,n){
		A[i]=readint();B[i]=readint();C[i]=readint();D[i]=readint();
	}
	int t1=0;
	REP(x,n){
		REP(y,n){
			AB[t1]=A[x]+B[y];
			CD[t1++]=C[x]+D[y];
		}
	}
	int ans=0;
	sort(AB,AB+t1);
	sort(CD,CD+t1);
	for(int i=0;i<t1;i++){
		ans+=upper_bound(CD,CD+t1,-AB[i])-lower_bound(CD,CD+t1,-AB[i]);
	}
	printf("%d\n",ans);
}