[BZOJ3926][Zjoi2015]诸神眷顾的幻想乡-后缀自动机

3926: [Zjoi2015]诸神眷顾的幻想乡

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

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

Sample Input

7 3

0 2 1 2 1 0 0

1 2

3 4

3 5

4 6

5 7

2 5

Sample Output

30

HINT

对于所有数据，1<=n<=100000, 1<=c<=10。
对于15%的数据，n<=2000。
另有5%的数据，所有空地都至多与两个空地相邻。
另有5%的数据，除一块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻。
另有5%的数据，除某两块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻

---

先放一张幽香sama的图片~~o((≧▽≦o)

咳咳，回归正文。
第一次做后缀自动机的题╮(╯▽╰)╭~
这题感觉……应该不算难的样子？？适合入门~
题意大概是求本质不同的串有多少个。
后缀自动机貌似刚好可以做到合并插入的相同子串~
所以写一个裸的后缀自动机。dfs插入，最后直接暴力统计答案。
所以咱这就做完了╮(╯▽╰)╭~

---

补充:
关于暴力统计答案的原理……
咱虽然明白，但在和一刚学后缀自动机的大触描述时(woc居然第一次就做这题,咱至少还打过一遍模板玩)却说得好乱……
隐约感觉这才是重点……

那就让咱描述一下吧:
对于后缀自动机这种神奇的东西，我们要明白它的本质~
比如这个fa指针，实际上它指向的是原串中当前节点所能代表的所有后缀，在集体缩短一段长度后得到的新的后缀集合中个数最少的，一个节点~

举个栗子:

komeijikoishi
  meijikoishi
    ijikoishi
         ishi

这些都是后缀。

比如咱匹配一个串komeijikomeji，
从头跑到第8个字符(‘o’)处时失配了，怎么办？
沿着fa跑，咱跑到了第2个字符(‘o’)和第8个字符(‘o’)处。
比较一下发现：

komeijikoishi
       koishi

这两个后缀的前缀相同(后缀的前缀即为子串),同时，我们的匹配长度变成了2。
这就是fa指针的本质，其实思路是很像咱的AC自动机的fail指针的，但它的一个节点代表的是多个后缀而不是一个~
(然而事实上AC自动机一个节点也能代表多个前缀)

那么回到正题~
本质不同的子串的统计方法就是：len[i]-len[fa[i]]……
为什么呢?
根据刚才的结论，有:
咱找到的当前节点由于失配往fa跳的缘故匹配长度少了len[i]-len[fa[i]],则根据上面的例子咱统计出了:

  omeijiko
   meijiko
    eijiko
     ijiko
      jiko
       iko
        ko

的情况~

注意到后缀自动机每个字符在每个len值下只会有一个节点。
而len值又表示了当前节点表示的所有后缀所能匹配到的最大长度。
fa指向的也是失配后变短最少的位置。
则我们可以保证每个子串仅会被统计一次。
所以这样直接统计是对的~

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<vector>
#include<algorithm>

using namespace std;

const int N=4000005;

inline int read()
{
    int x=0;
    char ch=getchar();
    while(ch<'0' || '9'<ch)ch=getchar();
    while('0'<=ch && ch<='9'){x=(x<<1)+(x<<3)+(ch^48);ch=getchar();}
    return x;
}

int n,c;
int color[100005],d[100005];
vector<int> g[100005];

struct SAM
{
    int pool,last;
    int fa[N],len[N],next[N][10];

    SAM(){pool=1;}

    int insert(int u,int v)
    {
        int now=++pool;
        len[now]=len[u]+1;

        while(!next[u][v] && u)
            next[u][v]=now,u=fa[u];

        if(!u)
            fa[now]=1;
        else
        {
            int q=next[u][v];
            if(len[q]==len[u]+1)
                fa[now]=q;
            else
            {
                int newq=++pool;
                fa[newq]=fa[q];
                len[newq]=len[u]+1;
                memcpy(next[newq],next[q],sizeof(next[q]));

                fa[now]=fa[q]=newq;

                while(next[u][v]==q)
                    next[u][v]=newq,u=fa[u];
            }
        }
        return now;
    }

    long long calc()
    {
        long long ans=0;
        for(int i=1;i<=pool;i++)
            ans+=(long long)len[i]-len[fa[i]];
        return ans;
    }
}koishi;

void dfs(int now,int fa,int pre)
{
    int npre=koishi.insert(pre,color[now]);
    for(int i=0;i<g[now].size();i++)
        if(g[now][i]!=fa)
            dfs(g[now][i],now,npre);
}

int main()
{
    n=read();
    c=read();
    for(int i=1;i<=n;i++)
        color[i]=read();

    for(int i=1;i<n;i++)
    {
        int u=read(),v=read();
        g[u].push_back(v);
        g[v].push_back(u);
        d[u]++;d[v]++;
    }

    for(int i=1;i<=n;i++)
        if(d[i]==1)dfs(i,0,1);

    printf("%lld\n",koishi.calc());

    return 0;
}