codechef Prime Distance On Tree（树分治+FFT）

Posted 2020-07-13 GFY

题目链接：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 }