太阳花田的结构比较特殊,只与一个空地相邻的空地数量不超过20个。
这是一棵只有20个叶子的树,容易想到对每个叶子节点dfs一遍,然后拉板子。。
其实就是搞一棵大trie树,求不同的子串个数,这个可以用SAM解决。
#include<iostream>
#include<cstdio>
#define N 200005
#define M 4000005
#define ll long long
using namespace std;
int n,C,a[N],x,y,q[M];ll Ans;
int F[M],st[M],c[M][15],cnt=1,vis[M];
int fst[N],nxt[N],to[N],l,d[N];
void link(int x,int y)
{
d[x]++;d[y]++;
to[++l]=y;nxt[l]=fst[x];fst[x]=l;
to[++l]=x;nxt[l]=fst[y];fst[y]=l;
}
int add(int p,int x)
{
int q,np,nq;
if (c[p][x]&&st[c[p][x]]==st[p]+1) return c[p][x];
st[np=++cnt]=st[p]+1;
while(p&&!c[p][x]) c[p][x]=np,p=F[p];
if (!p) F[np]=1;
else
{
if (st[p]+1==st[q=c[p][x]])
F[np]=q;
else
{
st[nq=++cnt]=st[p]+1;
for (int i=0;i<C;i++) c[nq][i]=c[q][i];
F[nq]=F[q];F[q]=F[np]=nq;
while(p&&c[p][x]==q)c[p][x]=nq,p=F[p];
}
}
return np;
}
void dfs(int x,int f,int b)
{
int t=add(b,a[x]);
for (int i=fst[x];i;i=nxt[i])
if (to[i]!=f) dfs(to[i],x,t);
}
int main()
{
scanf("%d%d",&n,&C);
for (int i=1;i<=n;i++)
scanf("%d",&a[i]);
for (int i=1;i<n;i++)
{
scanf("%d%d",&x,&y);link(x,y);
}
for (int i=1;i<=n;i++)
if (d[i]==1) dfs(i,0,1);
int l=0,r=1;
q[1]=1;
while(l<r)
{
x=q[++l];
if (x!=1) Ans+=st[x]-st[F[x]];
for (int i=0;i<C;i++)
if (c[x][i]&&!vis[c[x][i]]) vis[q[++r]=c[x][i]]=1;
}
printf("%lld\n",Ans);
}