[BZOJ3513] [MUTC2013] idiots - FFT快速傅里叶变换

3513: [MUTC2013]idiots

Time Limit: 20 Sec   Memory Limit: 128 MB
Submit: 373   Solved: 122
[ Submit][ Status][ Discuss]

Description

给定n个长度分别为a_i的木棒，问随机选择3个木棒能够拼成三角形的概率。

Input

第一行T(T<=100)，表示数据组数。
接下来若干行描述T组数据，每组数据第一行是n，接下来一行有n个数表示a_i。

Output

T行，每行一个整数，四舍五入保留7位小数。

Sample Input

2
4
1 3 3 4
4
2 3 3 4

Sample Output

0.5000000
1.0000000

HINT

T<=20

N<=100000

a_i<=100000 这是题目原本漏掉的条件！

Source

By sbullet

[ Submit][ Status][ Discuss] 
HOME   Back

        本来写的是NTT结果发现答案ans有可能大于模数，写longlonglong估计会TLE，所以就放弃了NTT写了一发FFT。

        突然感觉自己的代码好优美而且跑上了首页？（雾

        记得每次重做fill的时候一定要加上s[0]。

        好吧言归正传……

        首先我们看题目，既然a_i<=100000，我们考虑求出两根木棒的长度和为1-200000的方案数。那么怎么求出这个方案数呢？我们用A数组代表长度数组，B数组代表长度和数组，即B[x]代表两根木棒的长度和为x的方案数。

        那么我们可以得到B[i]=∑A[k]*A[i-k] 但是这时长度为p的木棒还和自身组成了长度为2*p的木棒，先不管。

        那么我们就可以FFT加速，因为B是卷积的形式，将A看成多项式即可。我们考虑补集，求不能拼成三角形的概率，再用1相减即可。

        因此可得ans=(C(n,3)-∑B[A[i]])/C(n,3)，其中C是组合数。

        再考虑之前和自身组成的木棒，只要扫一遍在B数组上直接减即可。

        时间复杂度O(Tnlogn)

#include "set"
#include "map"
#include "math.h"
#include "stdio.h"
#include "stdlib.h"
#include "iostream"
#include "algorithm"

using namespace std;

typedef long long ll;
const int N=300005,logn=20;
const double PI=acos(-1.0);
struct Val{
	double r,i;
	Val () {i=r=0.0;}
	Val (double _real,double _ima){
		r=_real; i=_ima;
	}
	Val operator+(const Val&x)const{
		return Val(r+x.r,i+x.i);
	}
	Val operator-(const Val&x)const{
		return Val(r-x.r,i-x.i);
	}
	Val operator*(const Val&x)const{
		return Val(r*x.r-i*x.i,i*x.r+r*x.i);
	}
} s[N];
int a[N],LEN,n;
ll sum[N],ans;

inline int read(){
	int v=0,f=1; char ch=getchar();
	while (ch<'0'||ch>'9'){ if(ch=='-') f=-1; ch=getchar();}
	while (ch<='9'&&ch>='0') { v=v*10+ch-'0'; ch=getchar();}
	return v;
}

inline void Rader(Val x[],int len){
	for (int i=1,j=len>>1,k;i<len-1;i++){
		if(i<j) swap(x[i],x[j]); k=len;
		while(j>=(k>>=1)) j-=k;
		if(j<=k) j+=k;
	}
}

inline void FFT(Val x[],int len,int c){
	Rader(x,len); int l,i,j,k; Val u,v;
	for (l=2;l<=len;l<<=1){ i=l>>1;
		Val w(cos(-c*2*PI/l),sin(-c*2*PI/l));
		for (j=0;j^len;j+=l){
			Val wn(1.0,0.0);
			for (k=j;k^(j+i);k++){
				u=x[k]; v=wn*x[k+i];
				x[k]=u+v; x[k+i]=u-v;
				wn=w*wn;
			}
		}
	}
	if (c==-1){
		for (i=0;i<len;i++){
			x[i].r/=len;
		}
	}
}

void init(){
	n=read(); int i,maxn=0;
	fill(s,s+LEN+1,Val(0.0,0.0));
	for(i=1;i<=n;i++)maxn=max(maxn,a[i]=read());
	for(LEN=1;(LEN>>1)<maxn;LEN<<=1);
	for(i=1;i<=n;i++)s[a[i]].r+=1.0;
}

void work(){
	FFT(s,LEN,1); int i; ans=0;
	for(i=0;i<LEN;i++){s[i]=s[i]*s[i];}
	FFT(s,LEN,-1); for(i=1;i<=n;i++)
	{s[2*a[i]].r-=1.0;} for(i=1;i<=LEN;i++)
	{sum[i]=sum[i-1]+(int)(s[i].r+0.5);}
	for(i=1;i<=n;i++)ans+=sum[a[i]];
	double p=ans*3.0/n/(n-1)/(n-2);
	printf("%.7lf\n",1.0-p);
}

int main(){
	int T; cin>>T;
	while(T--){
		init();
		work();
	}
	return 0;
}