[hdu 5307] He is Flying (FFT)

题意

给出长度为的数列，
对于任意,求出区间和为的区间的长度之和

题目分析

有了几道fft题的经验,套路就是把可加性的计算化为幂的次数做多项式卷积
有

分别把看成两个多项式，做FFT即可

• 拿举例

第一个多项式:
sum(i)的次数对应的系数加上i(1<=i<=n)
第二个多项式:
-sum(i)的次数对应的系数加上1(0<=i<=n-1)

由于负数次幂FFT无法处理,于是加上一个,即

第一个多项式:
sum(i)的次数对应的系数加上i(1<=i<=n)
第二个多项式:
s-sum(i)的次数对应的系数加上1(0<=i<=n-1)

所对应的答案存在的次数的系数上

• 可以类比
• 注意特判一下的情况
• 还有这道题卡精度,必须用Long Double,否则会WA(也可以写NTT)

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
using namespace std;
typedef long long LL;
typedef long double LD; //卡精度,不然会WA
const int MAXN = 1<<17;
const LD Pi = acos(-1.0);
struct complex
{
  LD r, i;
  complex(LD _r = 0, LD _i = 0):r(_r), i(_i){}

  complex operator +(const complex &t)const
  { return complex(r + t.r, i + t.i); }

  complex operator -(const complex &t)const
  { return complex(r - t.r, i - t.i); }
  
  complex operator *(const complex &t)const
  { return complex(r*t.r - i*t.i, i*t.r + t.i*r); }
  
  
}a1[MAXN], a2[MAXN];

inline void change(complex arr[], const int& len)
{
  for(register int i = 1, j = len/2; i < len-1; ++i)
  {
    if(i < j) swap(arr[i], arr[j]);
    int k = len/2;
    while(j >= k) j -= k, k>>=1;
    j += k;
  }
}

inline void fft(complex arr[], const int& len, const int& flg)
{
  change(arr, len);
  for(register int i = 2; i <= len; i<<=1)
  {
    complex wn = complex(cos(2*Pi*flg/i), sin(2*Pi*flg/i));
    for(register int j = 0; j < len; j += i)
    {
      complex w = complex(1, 0);
      for(register int k = j; k < j + i/2; ++k)
      {
        complex A0 = arr[k];
        complex wA1 = w * arr[k + i/2];
        arr[k] = A0 + wA1;
        arr[k + i/2] = A0 - wA1;
        w = w * wn;
      }
    }
  }
  if(flg == -1)
    for(register int i = 0; i < len; ++i)
      arr[i].r /= len;
}

int sum[MAXN], S;

LL ans[MAXN], pres[MAXN];

int main ()
{
  for(LL i = 1; i <= 100000; ++i) //预处理
    pres[i] = pres[i-1] + i*(i+1)/2;
  int T, n;
  scanf("%d", &T);
  while(T--)
  {
    scanf("%d", &n);
    int last = 0; LL ans0 = 0;
    for(int i = 1, x; i <= n; ++i)
    {
      scanf("%d", &x), sum[i] = sum[i-1] + x;
      if(!x) ++last;
      else ans0 += pres[last], last = 0;
    }
    printf("%lld\n", ans0 += pres[last]); //0特判
  
    S = sum[n];
  
    int len = 1;
    while(len < (S<<1)) len<<=1;
    
    for(int i = 0; i < len; ++i) a1[i] = a2[i] = complex(); //Part 1
    for(int i = 1; i <= n; ++i)
    {
      a1[sum[i]].r += i;
      if(i<n) a2[S-sum[i]].r += 1;
    }
    a2[S-0].r += 1; //i = sum[i] = 0
    fft(a1, len, 1); fft(a2, len, 1);
    for(int i = 0; i < len; ++i) a2[i] = a1[i] * a2[i];
    fft(a2, len, -1);
    for(int i = S+1; i <= (S<<1); ++i) ans[i-S] = LL(a2[i].r + 0.5);
    
    
    for(int i = 0; i < len; ++i) a1[i] = a2[i] = complex(); //Part 2
    for(int i = 1; i <= n; ++i)
    {
      a1[sum[i]].r += 1;
      if(i < n) a2[S-sum[i]].r += i;
    }
    fft(a1, len, 1); fft(a2, len, 1);
    for(int i = 0; i < len; ++i) a2[i] = a1[i] * a2[i];
    fft(a2, len, -1);
  
    for(int i = S+1; i <= (S<<1); ++i)
      printf("%lld\n", ans[i-S] - LL(a2[i].r + 0.5));
  }
}