【FFT求卷积】Problem D. Duel

  1 #include <stdio.h>
  2 #include <iostream>
  3 #include <string.h>
  4 #include <algorithm>
  5 #include <math.h>
  6 using namespace std;
  7 typedef long long ll;
  8 const double PI = acos(-1.0);
  9 struct complex
 10 {
 11     double r,i;
 12     complex(double _r = 0,double _i = 0)
 13     {
 14         r = _r; i = _i;
 15     }
 16     complex operator +(const complex &b)
 17     {
 18         return complex(r+b.r,i+b.i);
 19     }
 20     complex operator -(const complex &b)
 21     {
 22         return complex(r-b.r,i-b.i);
 23     }
 24     complex operator *(const complex &b)
 25     {
 26         return complex(r*b.r-i*b.i,r*b.i+i*b.r);
 27     }
 28 };
 29 void change(complex y[],int len)
 30 {
 31     int i,j,k;
 32     for(i = 1, j = len/2;i < len-1;i++)
 33     {
 34         if(i < j)swap(y[i],y[j]);
 35         k = len/2;
 36         while( j >= k)
 37         {
 38             j -= k;
 39             k /= 2;
 40         }
 41         if(j < k)j += k;
 42     }
 43 }
 44 void fft(complex y[],int len,int on)
 45 {
 46     change(y,len);
 47     for(int h = 2;h <= len;h <<= 1)
 48     {
 49         complex wn(cos(-on*2*PI/h),sin(-on*2*PI/h));
 50         for(int j = 0;j < len;j += h)
 51         {
 52             complex w(1,0);
 53             for(int k = j;k < j+h/2;k++)
 54             {
 55                 complex u = y[k];
 56                 complex t = w*y[k+h/2];
 57                 y[k] = u+t;
 58                 y[k+h/2] = u-t;
 59                 w = w*wn;
 60             }
 61         }
 62     }
 63     if(on == -1)
 64         for(int i = 0;i < len;i++)
 65             y[i].r /= len;
 66 }
 67 
 68 const int MAXN = 400040;
 69 complex x1[MAXN];
 70 int a[MAXN/4];
 71 long long num[MAXN];//100000*100000会超int
 72 char str[MAXN/4];
 73 int main()
 74 {
 75     freopen("duel.in","r",stdin);
 76     freopen("duel.out","w",stdout);
 77     int n=0;
 78     scanf("%s",str+1);
 79     int ls=strlen(str+1);
 80     for(int i=1;i<=ls;i++)
 81     {
 82         if(str[i]==‘1‘)
 83         {
 84             a[n++]=i;
 85         }
 86     }
 87     memset(num,0,sizeof(num));
 88     for(int i = 0;i < n;i++)
 89     {
 90         num[a[i]]++;
 91     }
 92     sort(a,a+n);
 93     int len1 = a[n-1]+1;
 94     int len = 1;
 95     while( len < 2*len1 )len <<= 1;
 96     for(int i = 0;i < len1;i++)
 97         x1[i] = complex(num[i],0);
 98     for(int i = len1;i < len;i++)
 99         x1[i] = complex(0,0);
100     fft(x1,len,1);
101     for(int i = 0;i < len;i++)
102         x1[i] = x1[i]*x1[i];
103     fft(x1,len,-1);
104     for(int i = 0;i < len;i++)
105         num[i] = (long long)(x1[i].r+0.5);
106     len = 2*a[n-1];
107     //减掉取两个相同的组合
108     for(int i = 0;i < n;i++)
109         num[a[i]+a[i]]--;
110     //选择的无序，除以2
111     for(int i = 1;i <= len;i++)
112     {
113         num[i]/=2;
114     }
115     ll ans=0;
116     for(int i=1;i<=len;i++)
117     {
118         if(i%2!=0) continue;
119         if(str[i/2]==‘0‘) continue;
120         ans+=num[i];
121     }
122     cout<<ans<<endl;
123     return 0;
124 }