[HDU 4609]3-idiots(FFT+生成函数)

题目链接

http://acm.hdu.edu.cn/showproblem.php?pid=4609

题目大意

给出$n$根木棍以及它们的长度，求取其中三根长度各不相同的木棍组成三角形的方案数。

思路

不妨设多项式$A(x)$,其中$x^i$前面的系数表示长度为$i$的木棍个数。那么多项式$A(x)^2$中$x^i$前面的系数表示用两根木棍拼成长度为$i$的木棍的方案数。

我们可以维护一个前缀和$sum[i]$表示对于所有的$j<=i$,$A(x)^2$多项式中$x^j$前面的系数，也就是用两根互不相同的木棍拼成长度小于等于$i$的方案数。这个前缀和可以在FFT做完以后完成。

然后我们枚举三根木棍中最长的那根木棍的长度$len$，求最长的木棍长度为$len$，三根木棍互不相同并能组成三角形的方案数。显然另外两根木棍的长度之和是大于$len$的，选择这两根木棍的方案数可以通过前缀和求补集得到(就是变成了后缀和)。然后减去这两根木棍中，一根木棍比$len$小，一根比$len$长的方案数，再减去两根都比$len$长的方案数，再减去有一根长度恰好为$len$的方案数，就得到了最长的木棍长度为$len$，三根木棍互不相同并能组成三角形的方案数。累加答案即可。

代码

#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <algorithm>
#include <complex>

#define MAXN 610000
#define PI 3.1415926535897384626

using namespace std;

typedef complex<long double>Complex;
typedef long long int LL;

inline void reverse(Complex a[],int n) //对长度为n的复数序列做FFT前必要的比特反转
{
    for(int i=1,j=n/2,k;i<n-1;i++)
    {
        if(i<j) swap(a[i],a[j]);
        k=n/2;
        while(j>=k)
        {
            j-=k;
            k>>=1;
        }
        if(j<k) j+=k;
    }
}

inline void FFT(Complex a[],int n,int flag) //对长度为n的复数序列a做FFT，flag=1是求值过程，flag=-1是插值过程
{
    reverse(a,n); //比特反转
    for(int i=1;i<n;i<<=1)
    {
        Complex wn=Complex(cos(PI/i),flag*sin(PI/i));
        for(int j=0;j<n;j+=(i<<1))
        {
            Complex w=Complex(1,0);
            for(int k=0;k<i;k++,w=w*wn)
            {
                Complex x=a[j+k],y=w*a[j+k+i];
                a[j+k]=x+y;
                a[j+k+i]=x-y; //!!!!
            }
        }
    }
    if(flag==-1) //插值操作还要乘上1/n的系数
        for(int i=0;i<n;i++)
            a[i]=Complex(a[i].real()/n,a[i].imag());
}

Complex a[MAXN];
LL len[MAXN]; //len[i]=第i根木棍的长度
LL num[MAXN]; //num[i]=2个木棍组合后长为i的方案数
LL sum[MAXN]; //sum[i]=2根不同的木棍组合后长度<=i的方案数

int main()
{
    int T;
    scanf("%d",&T);
    while(T--)
    {
        for(int i=0;i<MAXN;i++) a[i]=Complex(0,0);
        LL n=0,N=1,k=0; //n是序列长度，min{2^k}=min{N}>n
        scanf("%lld",&n);
        for(int i=0;i<n;i++)
        {
            scanf("%lld",&len[i]);
            a[len[i]]=Complex(a[len[i]].real()+1,a[len[i]].imag());
        }
        sort(len,len+n);
        int tmp=len[n-1]+1;
        for(;N<2*tmp+2;N<<=1,k++);
        FFT(a,N,1);
        for(int i=0;i<N;i++) a[i]=a[i]*a[i];
        FFT(a,N,-1);
        N=len[n-1]*2; //最高次项次数
        for(int i=0;i<=N;i++) num[i]=(LL)(a[i].real()+0.5);
        for(int i=0;i<n;i++) num[len[i]+len[i]]--; //去掉两个木棍相同的情况
        for(int i=0;i<=N;i++) num[i]/=2; //选择木棍时是无序的，而答案是有序的，除以2
        for(int i=1;i<=N;i++) sum[i]=sum[i-1]+num[i];
        LL ans=0; //ans=最终答案(选三根互不相同的棍组成一个三角形的方案数)
        for(int i=0;i<n;i++)
        {
            ans+=sum[N]-sum[len[i]]; //ans+=组合后的长度比木棍i长的两根不同的木棍的方案数
            ans-=(LL)(n-1-i)*i; //减去有一个取比木棍i长的，一个取比木棍i短的方案数
            ans-=n-1; //减去两根木棍中有一根是木棍i，另一根取其他木棍的方案数
            ans-=(LL)(n-1-i)*(n-i-2)/2; //!!!!!减去两根木棍都比木棍i长的方案数
        }
        LL all=(LL)n*(n-1)*(n-2)/6; //选三根互不相同的木棍总方案数
        printf("%.7lf\n",(double)ans/all);
    }
    return 0;
}