[bzoj3926][Zjoi2015]诸神眷顾的幻想乡

来自FallDream的博客，未经允许， 请勿转载，谢谢。

---

幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的2600岁生日，无数幽香的粉丝到了幽香家门前的太阳花田上来为幽香庆祝生日。 

粉丝们非常热情，自发组织表演了一系列节目给幽香看。幽香当然也非常高兴啦。 

这时幽香发现了一件非常有趣的事情，太阳花田有n块空地。在过去，幽香为了方便，在这n块空地之间修建了n-1条边将它们连通起来。也就是说，这n块空地形成了一个树的结构。 

有n个粉丝们来到了太阳花田上。为了表达对幽香生日的祝贺，他们选择了c中颜色的衣服，每种颜色恰好可以用一个0到c-1之间的整数来表示。并且每个人都站在一个空地上，每个空地上也只有一个人。这样整个太阳花田就花花绿绿了。幽香看到了，感觉也非常开心。 

粉丝们策划的一个节目是这样的，选中两个粉丝A和B（A和B可以相同），然后A所在的空地到B所在的空地的路径上的粉丝依次跳起来（包括端点），幽香就能看到一个长度为A到B之间路径上的所有粉丝的数目（包括A和B）的颜色序列。一开始大家打算让人一两个粉丝（注意：A,B和B,A是不同的，他们形成的序列刚好相反，比如红绿蓝和蓝绿红）都来一次，但是有人指出这样可能会出现一些一模一样的颜色序列，会导致审美疲劳。 

于是他们想要问题，在这个树上，一共有多少可能的不同的颜色序列（子串）幽香可以看到呢？ 

太阳花田的结构比较特殊，只与一个空地相邻的空地数量不超过20个。 

n<=100000 c<=10

因为叶子结点不超过20个，所以可以从每个叶子结点开始dfs，把节点插入同一个后缀自动机上。最后在后缀自动机上统计子串数量即可。

复杂度O(10*20*n)

#include<iostream>
#include<cstdio>
#include<algorithm>
#define MN 100000
#define ME 4000000
#define ll long long
using namespace std;
inline int read()
{
    int x = 0 , f = 1; char ch = getchar();
    while(ch < '0' || ch > '9'){ if(ch == '-') f = -1;  ch = getchar();}
    while(ch >= '0' && ch <= '9'){x = x * 10 + ch - '0';ch = getchar();}
    return x * f;
}

ll sum[ME+5];
int c[ME+5][11],n,m,step[ME+5],fail[ME+5],v[MN+5],sa[ME+5],in[MN+5],col[MN+5],cnt=0,head[MN+5];
struct edge{int to,next;}e[MN*2+5];
inline void ins(int f,int t)
{
    ++in[f];++in[t];
    e[++cnt]=(edge){t,head[f]};head[f]=cnt;
    e[++cnt]=(edge){f,head[t]};head[t]=cnt;
}

int Ins(int p,int x)
{
    int np=++cnt;step[np]=step[p]+1;
    for(;p&&!c[p][x];p=fail[p]) c[p][x]=np;
    if(!p) fail[np]=1;
    else
    {
        int q=c[p][x];
        if(step[q]==step[p]+1) fail[np]=q;
        else
        {
            int nq=++cnt;step[nq]=step[p]+1;
            fail[nq]=fail[q];fail[q]=fail[np]=nq;
            for(int i=0;i<m;++i) c[nq][i]=c[q][i];
            for(;c[p][x]==q;p=fail[p]) c[p][x]=nq;
        }
    }
    return np;
}

void Solve(int x,int fa,int p)
{
    int np=Ins(p,col[x]);
    for(int i=head[x];i;i=e[i].next)
        if(e[i].to!=fa) Solve(e[i].to,x,np);
}

int main()
{
    n=read();m=read();
    for(int i=1;i<=n;++i) col[i]=read();
    for(int i=1;i<n;++i) ins(read(),read());
    cnt=1;
    for(int i=1;i<=n;++i) if(in[i]==1) Solve(i,0,1);
    for(int i=1;i<=cnt;++i) ++v[step[i]];
    for(int i=1;i<=MN;++i) v[i]+=v[i-1];
    for(int i=1;i<=cnt;++i) sa[v[step[i]]--]=i;
    for(int i=cnt;i;--i)
    {
        sum[sa[i]]=sa[i]!=1;
        for(int j=0;j<m;++j) if(c[sa[i]][j])
            sum[sa[i]]+=sum[c[sa[i]][j]];
    }
    printf("%lld\n",sum[1]);
    return 0;
}