codechef Prime Distance On Tree（树分治+FFT）
题目链接：http://www.codechef.com/problems/PRIMEDST/
题意：给出一棵树，边长度都是1。每次任意取出两个点(u,v)，他们之间的长度为素数的概率为多大？
树分治，对于每个根出发记录边的长度出现几次，然后每次求卷积，用素数表查一下即可添加答案。
1 #include<algorithm> 2 #include<cstdio> 3 #include<cmath> 4 #include<cstring> 5 #include<iostream> 6 #include<complex> 7 #define ll long long 8 #define N 150005 9 typedef std::complex<double> cd; 10 int mark[200005],p[200005]; 11 int son[200005],F[200005],first[200005],next[200005],tot,go[200005]; 12 int lc,n,sum,dis[200005],root,vis[200005],mxdis; 13 ll a[200005],b[200005],c[200005],A[200005],B[200005]; 14 ll ans1=0,ans2=0; 15 int read(){ 16 char ch=getchar();int t=0,f=1; 17 while (ch<'0'||ch>'9'){ 18 if (ch=='-') f=-1; 19 ch=getchar(); 20 } 21 while ('0'<=ch&&ch<='9'){ 22 t=t*10+ch-'0'; 23 ch=getchar(); 24 } 25 return t*f; 26 } 27 int bitrev(int t,int n){ 28 int res=0; 29 for (int i=0;i<n;i++) res|=((t>>(n-i-1))&1)<<i;//括号要多加 30 return res; 31 } 32 void fft(cd *a,int n,int rev){ 33 int len=1<<n; 34 static cd y[N];double Pi=acos(-1); 35 for (int i=0;i<len;i++) y[i]=a[bitrev(i,n)]; 36 for (int d=1;d<len;d<<=1){ 37 cd wn(exp(cd(0,Pi*rev/d))); 38 for (int k=0;k<len;k+=2*d){ 39 cd w(1,0); 40 for (int i=k;i<k+d;i++,w*=wn){ 41 cd u=y[i],v=w*y[i+d]; 42 y[i]=u+v; 43 y[i+d]=u-v; 44 } 45 } 46 } 47 if (rev==-1) 48 for (int i=0;i<len;i++) y[i]/=len; 49 for (int i=0;i<len;i++) a[i]=y[i]; 50 } 51 void mul(ll *a,int la,ll *b,int lb,ll *c,int &lc){ 52 int len=1,n=0; 53 static cd t1[N],t2[N]; 54 for (;len<la*2||len<lb*2;len<<=1,n++); 55 for (int i=0;i<la;i++) t1[i]=cd(a[i],0); 56 for (int i=0;i<lb;i++) t2[i]=cd(b[i],0); 57 for (int i=la;i<len;i++) t1[i]=cd(0,0); 58 for (int i=lb;i<len;i++) t2[i]=cd(0,0); 59 fft(t1,n,1);fft(t2,n,1); 60 for (int i=0;i<len;i++) t1[i]*=t2[i]; 61 fft(t1,n,-1); 62 for (int i=0;i<len;i++) c[i]=(ll)(t1[i].real()+0.5); 63 lc=len; 64 } 65 void insert(int x,int y){ 66 tot++; 67 go[tot]=y; 68 next[tot]=first[x]; 69 first[x]=tot; 70 } 71 void add(int x,int y){ 72 insert(x,y); 73 insert(y,x); 74 } 75 void prework(){ 76 mark[0]=mark[1]=1; 77 for (int i=2;i<=50000;i++){ 78 if (!mark[i]){ 79 p[++p[0]]=i; 80 } 81 for (int j=1;j<=p[0]&&p[j]*i<=50000;j++){ 82 mark[p[j]*i]=1; 83 if (i%p[j]==0) break; 84 } 85 } 86 } 87 void findroot(int x,int fa){ 88 son[x]=1; 89 F[x]=0; 90 for (int i=first[x];i;i=next[i]){ 91 int pur=go[i]; 92 if (vis[pur]||pur==fa) continue; 93 findroot(pur,x); 94 son[x]+=son[pur]; 95 F[x]=std::max(F[x],son[pur]); 96 } 97 F[x]=std::max(F[x],sum-son[x]); 98 if (F[x]<F[root]) root=x; 99 } 100 void pre(int x,int fa){ 101 dis[x]=dis[fa]+1; 102 mxdis=std::max(mxdis,dis[x]); 103 for (int i=first[x];i;i=next[i]){ 104 int pur=go[i]; 105 if (pur==fa||vis[pur]) continue; 106 pre(pur,x); 107 } 108 } 109 void dfs(int x,int fa){ 110 b[dis[x]]++; 111 for (int i=first[x];i;i=next[i]){ 112 int pur=go[i]; 113 if (pur==fa||vis[pur]) continue; 114 dis[pur]=dis[x]+1; 115 dfs(pur,x); 116 } 117 } 118 void work(int x){ 119 dis[0]=-1;mxdis=0; 120 int all=sum; 121 pre(x,0); 122 for (int i=0;i<=mxdis+1;i++) a[i]=b[i]=0; 123 a[0]=1; 124 for (int i=first[x];i;i=next[i]){ 125 int pur=go[i]; 126 if (vis[pur]) continue; 127 dis[pur]=1; 128 for (int j=0;j<=mxdis;j++) b[j]=0; 129 dfs(pur,x); 130 for (int j=0;j<=mxdis;j++) A[j]=a[j],B[j]=b[j]; 131 mul(A,mxdis+1,B,mxdis+1,c,lc); 132 int lim=std::min(mxdis*2,50000); 133 for (int i=1;i<=p[0]&&p[i]<=lim;i++){ 134 ans1+=c[p[i]]; 135 } 136 for (int i=0;i<=mxdis;i++) 137 a[i]+=b[i]; 138 } 139 vis[x]=1; 140 for (int i=first[x];i;i=next[i]){ 141 int pur=go[i]; 142 if (vis[pur]) continue; 143 if (son[pur]>son[x]) sum=all-son[x]; 144 else sum=son[pur]; 145 root=0; 146 findroot(pur,x); 147 work(root); 148 } 149 } 150 int main(){ 151 prework(); 152 while (scanf("%d",&n)!=EOF){ 153 if (n==0) break; 154 tot=0; 155 for (int i=1;i<=n;i++) first[i]=vis[i]=0; 156 for (int i=1;i<n;i++){ 157 int x,y; 158 x=read();y=read(); 159 add(x,y); 160 } 161 root=0;sum=n; 162 F[0]=99999999; 163 ans1=ans2=0; 164 findroot(1,0); 165 work(root); 166 ans2=(ll)(((ll)n)*((ll)(n-1))/2); 167 double tmp=((double)(ans1))/((double)(ans2)); 168 printf("%.8lf\n",tmp); 169 } 170 }
