HDU 4609 3-idiots

FFT计算卷积

#include<bits/stdc++.h>
using namespace std;

const int MAXN=1<<18;
const long long mod=998244353;
const int G=3;
long long rev[MAXN],w[2][MAXN];
long long fac[MAXN],inv[MAXN];
long long A[MAXN];

long long powmod(long long x,long long p)
{
	long long ret=1;
	while(p)
	{
		if(p&1)
			ret=ret*x%mod;
		x=x*x%mod;
		p>>=1;
	}
	return ret;
}

void init()
{
	fac[0]=1;
	for(int i=1;i<MAXN;i++)
		fac[i]=fac[i-1]*i%mod;
	inv[MAXN-1]=powmod(fac[MAXN-1],mod-2);
	for(int i=MAXN-2;i>=0;i--)
		inv[i]=inv[i+1]*(i+1)%mod;
}

void initfft(int n)
{
	int i,k,x,y,v,dv;
	for(i=0;i<=n-1;i++)
	{
		x=i;
		y=0;
		for(k=1;k<n;k<<=1,x>>=1)
			(y<<=1)|=(x&1);
		rev[i]=y;
	}
	v=powmod(G,(mod-1)/n);
	dv=powmod(v,mod-2);
	w[0][0]=w[1][0]=1;
	for(i=1;i<=n-1;i++)
	{
		w[0][i]=w[0][i-1]*v%mod;
		w[1][i]=w[1][i-1]*dv%mod;
	}
}

inline void fft(long long *A,int n,int ff)
{
	int i,j,k,t,l,v,x,y;
	for(i=0;i<=n-1;i++)
	{
		if(i<rev[i])
			swap(A[i],A[rev[i]]);
	}
	for(i=1;i<n;i<<=1)
	{
		for(j=0,t=n/(i<<1);j<n;j+=(i<<1))
		{
			for(k=0,l=0;k<i;k++,l+=t)
			{
				x=A[i+j+k]*w[ff][l]%mod;
				y=A[j+k];
				A[j+k]=(x+y)%mod;
				A[i+j+k]=(y+mod-x)%mod;
			} 
		}
	}
	if(ff)
	{
		v=powmod(n,mod-2);
		for(i=0;i<=n-1;i++)
			A[i]=A[i]*v%mod;
	}
}

int num[MAXN];

int main()
{
	int T,n,i,m,mx,x;
	long long sum,tmp;
	double ans;
	scanf("%d",&T);
	while(T--)
	{
		scanf("%d",&n);
		memset(num,0,sizeof(num));
		mx=0;
		for(i=0;i<n;i++)
		{
			scanf("%d",&x);
			num[x]++;
			mx=max(mx,x);
		}
		m=1;
		while(m<=2*mx+2)
			m<<=1;
		initfft(m);
		for(i=0;i<m;i++)
			A[i]=num[i];
		fft(A,m,0);
		for(i=0;i<m;i++)
			A[i]=A[i]*A[i]%mod;
		fft(A,m,1);
		for(i=0;i<m;i++)
		{
			if(!(i&1))
				A[i]-=num[i>>1];
			A[i]>>=1;
		}
		tmp=0;
		sum=0;
		for(i=1;i<m;i++)
		{
			tmp+=A[i];
			sum+=tmp*num[i];
			
		}
		ans=1.0-sum*6.0/n/(n-1)/(n-2);
		printf("%.7f\n",ans);
	} 
}