1. 程式人生 > >The North American Invitational Programming Contest 2016:E. K-Inversions(FFT)

The North American Invitational Programming Contest 2016:E. K-Inversions(FFT)

You are given a string s consisting only of upper case letters A and B. For an integer k, a pair of indices i and j (1i<jn)(1 \le i < j \le n) is called a k-inversion if and only if s[i] = B , s[j] = A and ji=kj - i = k.

Consider the string BABA. It has two 1-inversions and one 3-inversion. It has no 2-inversions.

For each k between 1 and n - 1 (inclusive), print the number of k-inversions in the string s.

Input
Each input will consist of a single test case. Note that your program may be run multiple times on different inputs. The input will consist of a single line with a string s, which consists of only upper case As and Bs. The string s will be between 1 and 1,000,000 characters long. There will be no spaces.

Output
Output n−1 lines, each with a single integer. The first line’s integer should be the number of 1-inversions, the second should be the number of 2-inversions, and so on.

樣例輸入1
BABA
樣例輸出1
2
0
1
樣例輸入2
BBBBBAAAAA
樣例輸出2
1
2
3
4
5
4
3
2
1

思路:FFT。

我們可以把字串看成多項式。

把字元下標看成指數,即s[i]='A’對應的指數為i,s[j]='B’對應的指數為-j,則i-j即為對應的(i-j)-inversions。
然後可以把字串分為2個多項式,進行FFT,最後就可以得到答案。

#include<bits/stdc++.h>
using namespace std;
const int MAX=1e6+10;
const int MOD=1e9+7;
const double PI=acos(-1.0);
typedef long long ll;
struct Complex
{
    double re,im;
    Complex(double _re,double _im){re=_re;im=_im;}
    Complex(){}
}a[4*MAX],b[4*MAX];
Complex operator+(const Complex &c1,const Complex &c2){return Complex(c1.re+c2.re,c1.im+c2.im);}
Complex operator-(const Complex &c1,const Complex &c2){return Complex(c1.re-c2.re,c1.im-c2.im);}
Complex operator*(const Complex &c1,const Complex &c2){return Complex(c1.re*c2.re-c1.im*c2.im,c1.re*c2.im+c1.im*c2.re);}
int rev(int x,int len)
{
    int ret=0;
    for(int i=0;(1<<i)<len;i++)
    {
        ret<<=1;
        if(x&(1<<i))ret|=1;
    }
    return ret;
}
void FFT(Complex *p,int len,int DFT)
{
    for(int i=0;i<len;i++)
    {
        int x=rev(i,len);
        if(i<x)swap(p[i],p[x]);
    }
    for(int i=1;i<len;i*=2)
    {
        Complex wm=Complex(cos(DFT*PI/i),sin(DFT*PI/i));
        for(int k=0;k<len;k+=i*2)
        {
            Complex w=Complex(1,0);
            for(int j=0;j<i;j++)
            {
                Complex t=w*p[k+j+i];
                Complex u=p[k+j];
                p[k+j]=u+t;
                p[k+j+i]=u-t;
                w=w*wm;
            }
        }
    }
    if(DFT==-1)for(int i=0;i<len;i++)p[i].re/=len,p[i].im/=len;
}
char s[MAX];
int main()
{
    scanf("%s",s);
    int n=strlen(s);
    int len=0;
    while((1<<len)<=2*n)len++;
    len=1<<len;
    for(int i=0;i<n;i++)a[i]=Complex(s[i]=='A'?1:0,0);
    for(int i=0;i<n;i++)b[n-i-1]=Complex(s[i]=='B'?1:0,0);
    for(int i=n;i<len;i++)a[i]=b[i]=Complex(0,0);
    FFT(a,len,1);
    FFT(b,len,1);
    for(int i=0;i<len;i++)a[i]=a[i]*b[i];
    FFT(a,len,-1);
    for(int i=n;i<=2*n-2;i++)cout<<(int)(a[i].re+0.5)<<endl;
    return 0;
}