【ZJOI2015】bzoj3926 诸神眷顾的幻想乡

Description

幽香是全幻想乡里最受人欢迎的萌妹子，这天，是幽香的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个。 Input

第一行两个正整数n,c。表示空地数量和颜色数量。
第二行有n个0到c-1之间，由空格隔开的整数，依次表示第i块空地上的粉丝的衣服颜色。（这里我们按照节点标号从小到大的顺序依次给出每块空地上粉丝的衣服颜色）。
接下来n-1行，每行两个正整数u,v，表示有一条连接空地u和空地v的边。 Output

一行，输出一个整数，表示答案。

选取某一个点做根就可以看成一棵trie树，把trie建成SAM就可以求出在一个从根往下的路径上本质不同的子串个数。但是还有先上后下的子串无法处理。
因为叶子结点不多，我们可以以每个叶子为根建一棵trie，把这些trie建到一个SAM上。这样就可以涵盖所有的路径了。

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define LL long long
int rd()
{
    int x=0;
    char c=getchar();
    while (c<'0'||c>'9') c=getchar();
    while (c>='0'&&c<='9')
    {
        x=x*10+c-'0';
        c=getchar();
    }
    return x;
}
int fir[100010],ne[200010],to[200010],w[100010],fa[4000010],trans[4000010][12],val[4000010],n,c,tot=1;
void add(int num,int u,int v)
{
    ne[num]=fir[u];
    fir[u]=num;
    to[num]=v;
}
int add(int x,int p)
{
    int np,q,nq;
    val[np=++tot]=val[p]+1;
    while (p&&!trans[p][x])
    {
        trans[p][x]=np;
        p=fa[p];
    }
    if (!p) fa[np]=1;
    else
    {
        q=trans[p][x];
        if (val[q]==val[p]+1) fa[np]=q;
        else
        {
            val[nq=++tot]=val[p]+1;
            for (int i=0;i<c;i++) trans[nq][i]=trans[q][i];
            fa[nq]=fa[q];
            fa[q]=fa[np]=nq;
            while (p&&trans[p][x]==q)
            {
                trans[p][x]=nq;
                p=fa[p];
            }
        }
    }
    return np;
}
void dfs(int u,int fa,int last)
{
    int v,x=add(w[u],last);
    for (int i=fir[u];i;i=ne[i])
        if ((v=to[i])!=fa)
            dfs(v,u,x);
}
int main()
{
    //freopen("in","r",stdin);
    int u,v;
    LL ans=0;
    n=rd();
    c=rd();
    for (int i=1;i<=n;i++) w[i]=rd();
    for (int i=1;i<n;i++)
    {
        u=rd();
        v=rd();
        add(i<<1,u,v);
        add(i<<1|1,v,u);
    }
    for (int i=1;i<=n;i++)
        if (!ne[fir[i]])
        {
            u=add(w[i],1);
            for (int j=fir[i];j;j=ne[j]) dfs(to[j],i,u);
        }
    for (int i=1;i<=tot;i++) ans+=val[i]-val[fa[i]];
    printf("%lld\n",ans);
}