CF 528D. Fuzzy Search NTT

CF 528D. Fuzzy Search NTT

题目大意

给出文本串S和模式串T和k，S,T为DNA序列(只含ATGC)。对于S中的每个位置\(i\)，只要中[i-k,i+k]有一个位置匹配了字符\(i\)，那么就认为\(i\)可以匹配。求S中有多少位置匹配了T。

思路

一共有四个字母，我们分别计算每个字母是否可行，其他不管。

最后四个都满足的位置就是一个合法位置(指的是初始位置)。

设g[i]表示S_i位置是否是枚举的字母，f[i]表示M_i是否是是枚举的字母。

他们满足条件只需要右斜对角线==len

发现每个点又是右斜对角线,反转ntt

错误

有点zz，4写成了m，还忘记删掉调试了，wrong了两发，ntt真好调试(不用调试)。

代码

#include <bits/stdc++.h>
using namespace std;
const int N=1e6+7,mod=998244353;
int read() {
	int x=0,f=1;char s=getchar();
	for(;s>'9'||s<'0';s=getchar()) if(s=='-') f=-1;
	for(;s>='0'&&s<='9';s=getchar()) x=x*10+s-'0';
	return x*f;
}
int n,m,k,limit=1,l,r[N];
char S[N],T[N];
int q_pow(int a,int b) {
	int ans=1;
	while(b) {
		if(b&1) ans=1LL*ans*a%mod;
		a=1LL*a*a%mod;
		b>>=1;
	}
	return ans;
}
void ntt(int *a,int type) {
	for(int i=0;i<=limit;++i)
		if(i<r[i]) swap(a[i],a[r[i]]);
	for(int mid=1;mid<limit;mid<<=1) {
		int Wn=q_pow(3,(mod-1)/(mid<<1));
		for(int i=0;i<limit;i+=(mid<<1)) {
			for(int j=0,w=1;j<mid;j++,w=1LL*w*Wn%mod) {
				int x=a[i+j],y=1LL*w*a[i+j+mid]%mod;
				a[i+j]=(x+y)%mod;
				a[i+j+mid]=(x+mod-y)%mod;
			}
		}
	}
	if(type==-1) {
		reverse(&a[1],&a[limit]);
		int inv=q_pow(limit,mod-2);
		for(int i=0;i<=limit;++i) a[i]=1LL*a[i]*inv%mod;
	}
}
int AAA[N],f[N],g[N],sum[N];
void solve(char x) {
	memset(f,0,sizeof(f));
	memset(g,0,sizeof(g));
	int tong,h=0,d=0;
	sum[0]=(S[0]==x);
	for(int i=1;i<n;++i) sum[i]=sum[i-1]+(S[i]==x);
	for(int i=0;i<n;++i) {
		int y=i+k>=n ? sum[n-1] : sum[i+k];
		int x=i-k-1 < 0 ? 0 : sum[i-k-1];
		g[i]=(bool)(y-x);
	}
	for(int i=0;i<m;++i) f[m-i-1]=(T[i]==x);
	// for(int i=0;i<n;++i) cout<<g[i]<<" ";cout<<"\n";
	// for(int i=0;i<m;++i) cout<<f[i]<<" ";cout<<"\n";

	ntt(g,1),ntt(f,1);
	for(int i=0;i<=limit;++i) f[i]=1LL*f[i]*g[i]%mod;
	ntt(f,-1);
	int gs=0;
	for(int i=0;i<m;++i) gs+=(T[i]==x);
	for(int i=m-1,js=0;i<=n-1;++i,++js) AAA[js]+=(f[i]==gs);
}
int main() {
	n=read(),m=read(),k=read();
	scanf("%s%s",S,T);
	while(limit<=n+m-2) limit<<=1,l++;
	for(int i=0;i<=limit;++i)
		r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));
	solve('A'),solve('T'),solve('G'),solve('C');
	int tot=0;
	for(int i=0;i<=n-m+1;++i) tot+=(AAA[i]==4);
	printf("%d\n",tot);
	return 0;
}