【bzoj3926】[Zjoi20150]诸神眷顾的幻想乡 后缀自动机+trie

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%的数据，除某两块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻

Source

OrzCLJ

因为叶子只有二十个，所以以这些叶子为根建trie树，然后在树上建后缀自动机。

p状态代表的串有p -> val个，其中p -> par -> val个重复，答案累加p -> val - p -> par -> val即可。

dfs建后缀自动机，太神啦！

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<iostream>
using namespace std;

typedef long long LL;
const int SZ = 4000010;

int n,c;

struct sam_node{
    sam_node *par,*ch[10];
    int val;
}T[SZ], *root;

int Tcnt = 0;

sam_node* newnode(int x)
{
    sam_node *k = T + (Tcnt ++);
    k -> par = 0;
    memset(k -> ch,0,sizeof(k -> ch));
    k -> val = x;
    return k;
}

sam_node* insert(sam_node *last,int x)
{
    sam_node *p = last;
    sam_node *np = newnode(last -> val + 1);
    while(p && !p -> ch[x])
        p -> ch[x] = np,p = p -> par;
    if(!p)
        np -> par = root;
    else
    {
        sam_node *q = p -> ch[x];
        if(q -> val == p -> val + 1)
            np -> par = q;
        else
        {
            sam_node *nq = newnode(p -> val + 1);
            memcpy(nq -> ch,q -> ch,sizeof(nq -> ch));
            nq -> par = q -> par;
            q -> par = np -> par = nq;
            while(p && p -> ch[x] == q)
                p -> ch[x] = nq,p = p -> par;
        }
    }
    return np;
}

int head[SZ],nxt[SZ],to[SZ],tot = 0;

void build(int f,int t)
{
    to[++ tot] = t;
    nxt[tot] = head[f];
    head[f] = tot;
}


bool vis[SZ];
int C[SZ],in[SZ];

void dfs(sam_node *p,int u,int x,int fa)
{
    sam_node *now = insert(p,x);
    for(int i = head[u];i;i = nxt[i])
    {
        int v = to[i];
        if(v == fa) continue;
        dfs(now,v,C[v],u);
    }
}

void init()
{
    root = newnode(0);
}

int main()
{
    init();
    scanf("%d%d",&n,&c);
    for(int i = 1;i <= n;i ++) scanf("%d",&C[i]);
    for(int i = 1;i < n;i ++)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        build(u,v);
        build(v,u);
        in[u] ++; in[v] ++;
    }

    for(int i = 1;i <= n;i ++)
        if(in[i] == 1)
            dfs(root,i,C[i],0);
    LL ans = 0;
    for(int i = 1;i < Tcnt;i ++)
    {
        sam_node *p = T + i;
        ans += p -> val - p -> par -> val;
    }
    printf("%lld",ans);
    return 0;
}