[BZOJ1921] [CTSC2010]珠宝商

Description

Input

第一行包含两个整数 N,M，表示城市个数及特征项链的长度。 接下来的N-1 行， 每行两个整数 x,y， 表示城市 x 与城市 y 有直接道路相连。城市由1~N进行编号。接下来的一行，包含一个长度为 N，仅包含小写字母的字符串，第 i 位的字符表示在城市 i 流行的原料类型。 最后一行， 包含一个长度为 M， 仅包含小写字母的字符串， 表示特征字符串。

Output

仅包含一个整数,为 N2 * Expectation

Sample Input

3 5 
1 2 
1 3 
aab 
abaab  

Sample Output

15 

Solution

点分治+后缀自动机....

思路是真的神奇...看了好久题解才看明白。

先考虑暴力，有一个很显然的\(O(n^2)\)的暴力：

• 枚举每个点作为起点，\(dfs\)另一个点，\(dfs\)的同时在特征字符串\(S\)的\(SAM\)上跑，顺便统计答案就好了。
• 这个在\(SAM\)上跑实际上就相当于每次在一个已经匹配了的串后面加一个字符，那么直接沿着\(SAM\)的转换边走就好了。

换个角度思考，还有另一种暴力：

• 枚举一个点\(x\)，统计出每条以这个点为\(lca\)的路径的代价。
• 考虑这个看起来高级一点的暴力怎么做，
• 对于点\(x\)，路径肯定是\(a\to x \to b\)的形式，那么我们把他拆成两段\(a\to x\)和\(x \to b\)，注意到如果我们把特征字符串反过来，那么这两种其实就是一样的，所以我们现在考虑\(a \to x\)怎么统计，然后在翻转过的\(S\)上再做一遍就好了。
• 那么我们就是要统计出对于自动机上的点\(i\)，有多少以\(x\)结尾的路径字符串在这个点。
• 那么问题就相当于当前已经匹配了一个串，要在这个串前面加一个字符，那么沿着\(parent\)树跳就好了。
• 最后答案就是对于字符串上每个点，这个点开始和这个点结束的方案之积 的和。

注意到第二种暴力可以用点分治优化，那么第二种暴力的复杂度就是\(O(n\log n+nm)\)。

考虑如何把后面一项优化一下，注意到第一种方案是不需要每次都扫一遍自动机的，所以可以在点分治的时候设一个阀值\(B\)，若\(size>B\)就用第二种，否则用第一种。

可以发现当\(B=\sqrt{n}\)的时候复杂度最优，此时时间复杂度为\(O((n+m)\sqrt{n})\)。

#include<bits/stdc++.h>
using namespace std;

#define ll long long 

void read(int &x) {
    x=0;int f=1;char ch=getchar();
    for(;!isdigit(ch);ch=getchar()) if(ch=='-') f=-f;
    for(;isdigit(ch);ch=getchar()) x=x*10+ch-'0';x*=f;
}
 
void print(ll x) {
    if(x<0) putchar('-'),x=-x;
    if(!x) return ;print(x/10),putchar(x%10+48);
}
void write(ll x) {if(!x) putchar('0');else print(x);putchar('\n');}

const int maxn = 2e5+10;
const int inf = 1e9;

ll ans;
char s[maxn];
int n,m,rt,siz,B,top;
int head[maxn],tot,a[maxn];
int vis[maxn],f[maxn],sz[maxn],tmp[maxn],t[maxn];
struct edge{int to,nxt;}e[maxn<<1];

void ins(int u,int v) {e[++tot]=(edge){v,head[u]},head[u]=tot;}

struct Suffix_Automaton {
	int cnt,qs,lstp;
	int par[maxn],tr[maxn][26],ml[maxn],pos[maxn],sz[maxn];
	int t[maxn],r[maxn],son[maxn][26],str[maxn],tag[maxn],rev[maxn];

	void append(int x,int v) {
		int p=lstp,np=++cnt;pos[np]=v,sz[np]=1,ml[np]=ml[p]+1,rev[v]=np;lstp=np;
		for(;p&&tr[p][x]==0;p=par[p]) tr[p][x]=np;
		if(!p) return par[np]=qs,void();
		int q=tr[p][x];
		if(ml[p]+1<ml[q]) {
			int nq=++cnt;ml[nq]=ml[p]+1;
			memcpy(tr[nq],tr[q],sizeof tr[nq]);
			par[nq]=par[q],par[q]=par[np]=nq;
			for(;p&&tr[p][x]==q;p=par[p]) tr[p][x]=nq;
		} else par[np]=q;
	}

	void prepare(char *ss) {
		lstp=qs=cnt=1;
		for(int i=1;i<=m;i++) append(str[i]=ss[i]-'a',i);

		for(int i=1;i<=cnt;i++) t[ml[i]]++;
		for(int i=1;i<=m;i++) t[i]+=t[i-1];
		for(int i=1;i<=cnt;i++) r[t[ml[i]]--]=i;

		for(int i=cnt;i;i--) {
			int p=r[i];
			sz[par[p]]+=sz[p];
			if(!pos[par[p]]) pos[par[p]]=pos[p];
			son[par[p]][str[pos[p]-ml[par[p]]]]=p;
		}
	}

	void mark(int x,int fa,int now,int len) {
		if(len==ml[now]) now=son[now][a[x]];
		else if(str[pos[now]-len]!=a[x]) now=0;
		if(!now) return ;len++;tag[now]++;
		for(int i=head[x];i;i=e[i].nxt)
			if(e[i].to!=fa&&!vis[e[i].to]) mark(e[i].to,x,now,len);
	}

	void push() {for(int i=1;i<=cnt;i++) tag[r[i]]+=tag[par[r[i]]];}
}sam1,sam2;

void get_rt(int x,int fa) {
	sz[x]=1,f[x]=0;
	for(int i=head[x];i;i=e[i].nxt)
		if(e[i].to!=fa&&!vis[e[i].to]) 
			get_rt(e[i].to,x),sz[x]+=sz[e[i].to],f[x]=max(f[x],sz[e[i].to]);
	f[x]=max(f[x],siz-sz[x]);
	if(f[x]<f[rt]) rt=x;
}

void get_node(int x,int fa) {
	t[++top]=x;
	for(int i=head[x];i;i=e[i].nxt)
		if(e[i].to!=fa&&!vis[e[i].to]) get_node(e[i].to,x);
}

void dfs(int x,int fa,int now) {
	now=sam1.tr[now][a[x]];
	if(!now) return ;
	ans+=sam1.sz[now];
	for(int i=head[x];i;i=e[i].nxt)
		if(e[i].to!=fa&&!vis[e[i].to]) dfs(e[i].to,x,now);
}

void work(int x,int fa,int op) {
	memset(sam1.tag,0,(sam1.cnt+2)*4);
	memset(sam2.tag,0,(sam2.cnt+2)*4);
	if(fa) sam1.mark(x,fa,sam1.tr[1][a[fa]],1),sam2.mark(x,fa,sam2.tr[1][a[fa]],1);
	else sam1.mark(x,fa,1,0),sam2.mark(x,fa,1,0);
	sam1.push(),sam2.push();
	for(int i=1;i<=m;i++) ans+=1ll*op*sam1.tag[sam1.rev[i]]*sam2.tag[sam2.rev[m-i+1]];
}

void solve(int x) {
	get_rt(x,0);siz=sz[x];
	if(siz<=B) {
		top=0,get_node(x,0);
		for(int i=1;i<=top;i++) dfs(t[i],0,sam1.qs);
		for(int i=1;i<=top;i++) vis[t[i]]=0;
		return ;
	}
	for(int i=head[x];i;i=e[i].nxt) tmp[e[i].to]=sz[e[i].to];
	work(x,0,1);
	for(int i=head[x];i;i=e[i].nxt) if(!vis[e[i].to]) work(e[i].to,x,-1);
	vis[x]=1;
	for(int i=head[x];i;i=e[i].nxt)
		if(!vis[e[i].to]) siz=tmp[e[i].to],rt=0,get_rt(e[i].to,x),solve(rt);
}

int main() {
	read(n),read(m);B=sqrt(n);
	for(int i=1,x,y;i<n;i++) read(x),read(y),ins(x,y),ins(y,x);
	scanf("%s",s+1);
	for(int i=1;i<=n;i++) a[i]=s[i]-'a';
	scanf("%s",s+1);
	sam1.prepare(s);
	reverse(s+1,s+m+1);
	sam2.prepare(s);
	siz=n,f[0]=inf,get_rt(1,0),solve(rt);write(ans);
	return 0;
}