【BZOJ4598】[Sdoi2016]模式字符串 树分治+hash

【BZOJ4598】[Sdoi2016]模式字符串

Description

给出n个结点的树结构T，其中每一个结点上有一个字符，这里我们所说的字符只考虑大写字母A到Z,再给出长度为m的模式串s，其中每一位仍然是A到z的大写字母。Alice希望知道，有多少对结点<u，v>满足T上从u到V的最短路径形成的字符串可以由模式串S重复若干次得到？这里结点对<u，v>是有序的，也就是说<u，v>和<v，u>需要被区分.

所谓模式串的重复，是将若干个模式串S依次相接（不能重叠).例如当S=PLUS的时候，重复两次会得到PLUSPLUS，重复三次会得到PLUSPLUSPLUS,同时要注恿，重复必须是整数次的。例如当S=XYXY时，因为必须重复整数次，所以XYXYXY不能看作是S重复若干次得到的。

Input

每一个数据有多组测试，

第一行输入一个整数C，表示总的测试个数。

对于每一组测试来说：

第一行输入两个整数，分别表示树T的结点个数n与模式长度m。结点被依次编号为1到n，

之后一行，依次给出了n个大写字母（以一个长度为n的字符串的形式给出），依次对应树上每一个结点上的字符（第i个字符对应了第i个结点).

之后n-1行，每行有两个整数u和v表示树上的一条无向边，之后一行给定一个长度为m的由大写字母组成的字符串，为模式串S。

1<=C<=10,3<=N<=10000003<=M<=1000000

Output

给出C行，对应C组测试。每一行输出一个整数，表示有多少对节点<u,v>满足从u到v的路径形成的字符串恰好是模式串的若干次重复.

Sample Input

1
11 4
IODSSDSOIOI
1 2
2 3
3 4
1 5
5 6
6 7
3 8
8 9
6 10
10 11
SDOI

Sample Output

5

题解：考虑点分治。对于以x为根的点分治结构中，令f[i]表示所有深度%m=i的点中，有多少y满足y到x的路径是m的前i位。g[i]表示所有深度%m=i的点中，有多少y满足x到y的路径是m的后i位。然后采用树形DP式的方法统计答案，用hash来判断一个串是否是模式串的一部分即可。

#include <cstdio>
#include <cstring>
#include <iostream>
using namespace std;
const int maxn=1000010;
typedef long long ll;
typedef unsigned long long ull;
int cas,n,m,cnt,tot,mn,rt;
ll ans;
int to[maxn<<1],next[maxn<<1],head[maxn],siz[maxn],vis[maxn],f[maxn],g[maxn],sf[maxn],sg[maxn];
ull h1[maxn],h2[maxn],bs[maxn];
char S[maxn],T[maxn];
void getrt(int x,int fa)
{
	siz[x]=1;
	int tmp=0,i;
	for(i=head[x];i!=-1;i=next[i])	if(!vis[to[i]]&&to[i]!=fa)
		getrt(to[i],x),siz[x]+=siz[to[i]],tmp=max(tmp,siz[to[i]]);
	tmp=max(tmp,tot-siz[x]);
	if(tmp<mn)	mn=tmp,rt=x;
}
int getdep(int x,int fa,int dep,ull hs)
{
	siz[x]=1,hs=hs*1331+S[x];
	if(h1[dep]==hs)	f[(dep-1)%m+1]++,ans+=sg[m-(dep-1)%m];
	if(h2[dep]==hs)	g[(dep-1)%m+1]++,ans+=sf[m-(dep-1)%m];
	int i,tmp=1;
	for(i=head[x];i!=-1;i=next[i])	if(!vis[to[i]]&&to[i]!=fa)
		siz[x]+=siz[to[i]],tmp=max(tmp,getdep(to[i],x,dep+1,hs)+1);
	return tmp;
}
void dfs(int x)
{
	vis[x]=1;
	int i,j,k,tmp=0;
	sf[1]=sg[1]=1;
	for(i=head[x];i!=-1;i=next[i])	if(!vis[to[i]])
	{
		k=min(m,getdep(to[i],x,2,S[x])+1),tmp=max(tmp,k);
		for(j=1;j<=k;j++)	sf[j]+=f[j],sg[j]+=g[j],f[j]=g[j]=0;
	}
	for(i=1;i<=tmp;i++)	sf[i]=sg[i]=0;
	for(i=head[x];i!=-1;i=next[i])	if(!vis[to[i]])	tot=siz[to[i]],mn=1<<30,getrt(to[i],x),dfs(rt);
}
inline int rd()
{
	int ret=0,f=1;	char gc=getchar();
	while(gc<'0'||gc>'9')	{if(gc=='-')	f=-f;	gc=getchar();}
	while(gc>='0'&&gc<='9')	ret=ret*10+(gc^'0'),gc=getchar();
	return ret*f;
}
inline void add(int a,int b)
{
	to[cnt]=b,next[cnt]=head[a],head[a]=cnt++;
}
void work()
{
	n=rd(),m=rd(),ans=0;
	scanf("%s",S+1);
	int i,a,b;
	memset(head,-1,sizeof(head[0])*(n+1)),cnt=0,memset(vis,0,sizeof(vis[0])*(n+1));
	for(i=1;i<n;i++)	a=rd(),b=rd(),add(a,b),add(b,a);
	scanf("%s",T+1);
	for(bs[0]=1,i=1;i<=n;i++)	bs[i]=bs[i-1]*1331,h1[i]=h1[i-1]+T[(i-1)%m+1]*bs[i-1],h2[i]=h2[i-1]+T[m-(i-1)%m]*bs[i-1];
	tot=n,mn=1<<30,getrt(1,0),dfs(rt);
	printf("%lld\n",ans);
}
int main()
{
	cas=rd();
	while(cas--)	work();
	return 0;
}