[Codeforces 580D]Fizzy Search(FFT)

题面

给定母串和模式串，字符集大小为4，给定k，模式串在某个位置匹配当且仅当任意位置模式串的这个字符所对应的母串的位置的左右k个字符之内有一个与它相同的，求模式串能全部匹配的次数。

分析

我们先考虑\(k=0\)的情况,即一般的字符串匹配。设母串为\(S\),模式串为\(T\),\(ans_i\)表示母串从\(i\)位置开始与\(T\)匹配,能够匹配的字符个数(注意:当遇到不匹配的字符时仍继续匹配，直到匹配完整个串)

\[ans_p=\sum_{i=1}^m [S_{p+i-1}=T_i]\]

注意到\((p+i-1)+i\)不是常数，不符合卷积的形式。令\(T_i=T_{m-i+1}\),则

\[ans_p=\sum_{i=1}^m [S_{p+i-1}=T_{m-i+1}]\]

这样\((p+i-1)+(m-i+p)=m+p\)为常数，符合卷积的形式。但是现在仍然无法FFT处理。

容易发现，每个字符的贡献(即这个字符的匹配个数)是可加的。那么我们可以枚举字符\(c\),设\(a_{i-1}=[S_i=c],b_{m-i}=[T_i=c]\),这样\(a\)和\(b\)卷积时只有两个位置都为1的时候匹配，对答案产生1的贡献。因此\(ans_i+=(a \cdot b)_i\).枚举完字符后，只需要遍历\(ans\)序列,如果\(ans_i=m\),则说明该位置能够与\(T\)匹配

对于\(k>0\)的情况，我们只需要稍加修改\(a\)的定义。若\([i-k,i+k]\)中存在字符\(c\),则我们令\(a_{i-1}=1\),否则为0. 可以预处理前缀和来判断。这样就可以FFT了

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath> 
#define maxn 1048576
#define maxc 4
using namespace std;
const double pi=acos(-1.0);
struct com{
    double real;
    double imag;
    com(){
        
    } 
    com(double _real,double _imag){
        real=_real;
        imag=_imag;
    }
    com(double x){
        real=x;
        imag=0;
    }
    void operator = (const com x){
        this->real=x.real;
        this->imag=x.imag;
    }
    void operator = (const double x){
        this->real=x;
        this->imag=0;
    }
    friend com operator + (com p,com q){
        return com(p.real+q.real,p.imag+q.imag);
    }
    friend com operator + (com p,double q){
        return com(p.real+q,p.imag);
    }
    void operator += (com q){
        *this=*this+q;
    }
    void operator += (double q){
        *this=*this+q;
    }
    friend com operator - (com p,com q){
        return com(p.real-q.real,p.imag-q.imag);
    }
    friend com operator - (com p,double q){
        return com(p.real-q,p.imag);
    }
    void operator -= (com q){
        *this=*this-q;
    }
    void operator -= (double q){
        *this=*this-q;
    }
    friend com operator * (com p,com q){
        return com(p.real*q.real-p.imag*q.imag,p.real*q.imag+p.imag*q.real);
    }
    friend com operator * (com p,double q){
        return com(p.real*q,p.imag*q);
    } 
    void operator *= (com q){
        *this=(*this)*q;
    }
    void operator *= (double q){
        *this=(*this)*q;
    }
    friend com operator / (com p,double q){
        return com(p.real/q,p.imag/q);
    } 
    void operator /= (double q){
        *this=(*this)/q;
    } 
    void print(){
        printf("%lf + %lf i ",real,imag);
    }
};
int rev[maxn+5];
void fft(com *x,int n,int type){
    for(int i=0;i<n;i++) if(i<rev[i]) swap(x[i],x[rev[i]]);
    for(int len=1;len<n;len*=2){
        int sz=len*2;
        com wn1=com(cos(2*pi/sz),type*sin(2*pi/sz));
        for(int l=0;l<n;l+=sz){
            int r=l+len-1;
            com wnk=1;
            for(int i=l;i<=r;i++){
                com tmp=x[i+len];
                x[i+len]=x[i]-wnk*tmp;
                x[i]=x[i]+wnk*tmp;
                wnk=wnk*wn1;
            }
        }
    }
    if(type==-1) for(int i=0;i<n;i++) x[i]/=n;
} 

inline int get_id(char c){
    if(c=='A') return 0;
    else if(c=='T') return 1;
    else if(c=='G') return 2;
    else return 3;
}
int n,m,K;
char s[maxn+5],t[maxn+5];
int sum[maxc+5][maxn+5];
int match[maxn+5][maxc+5];//标记s的第i位周围有没有字符j 

com a[maxn+5],b[maxn+5];
long long ans[maxn+5];
int main(){
    scanf("%d %d %d",&n,&m,&K);
    scanf("%s",s+1);
    scanf("%s",t+1);
    for(int i=1;i<=n;i++){
        for(int j=0;j<maxc;j++) sum[j][i]=sum[j][i-1]+(get_id(s[i])==j);
    } 
    for(int i=1;i<=n;i++){
        int lb=max(i-K,1);
        int rb=min(i+K,n);
        for(int j=0;j<maxc;j++){
            if(sum[j][rb]-sum[j][lb-1]>0) match[i][j]=1;
            else match[i][j]=0;
        }
    }
    int M=n+m;
    int N=1,L=0;
    while(N<=M){
        N*=2;
        L++;
    }
    for(int i=0;i<N;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
    for(int c=0;c<maxc;c++){
        memset(a,0,sizeof(a));
        memset(b,0,sizeof(b));
        for(int i=1;i<=n;i++){
            if(match[i][c]) a[i-1]=1;
            else a[i-1]=0;
        }
        for(int i=1;i<=m;i++){
            if(get_id(t[i])==c) b[m-i]=1;
            else b[m-i]=0;
        }
        fft(a,N,1);
        fft(b,N,1);
        for(int i=0;i<N;i++) a[i]*=b[i];
        fft(a,N,-1);
        for(int i=0;i<N;i++) ans[i]+=(long long)(a[i].real+0.5);
    } 
    int cnt=0;
    for(int i=0;i<N;i++) if(ans[i]==m) cnt++;
    printf("%d\n",cnt);
}