诸神眷顾的幻想乡

Posted 2020-10-09 mxzf0213

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

Time Limit: 10 Sec  Memory Limit: 512 MB

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

分析：求树上不同路径的个数；

　　　把每个叶子提起来作为树根，可以发现所有子串必定会出现在到根的路径上；

　　　所以以每个叶子为根dfs，作一个广义后缀自动机；

　　　最后只需求出每个节点产生的新子串个数即可；

代码：

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <algorithm>
#include <climits>
#include <cstring>
#include <string>
#include <set>
#include <bitset>
#include <map>
#include <queue>
#include <stack>
#include <vector>
#include <cassert>
#include <ctime>
#define rep(i,m,n) for(i=m;i<=(int)n;i++)
#define inf 0x3f3f3f3f
#define mod 1000000007
#define vi vector<int>
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define ll long long
#define pi acos(-1.0)
#define pii pair<int,int>
#define sys system("pause")
#define ls (rt<<1)
#define rs (rt<<1|1)
#define all(x) x.begin(),x.end()
const int maxn=1e5+10;
const int N=5e4+10;
using namespace std;
ll gcd(ll p,ll q){return q==0?p:gcd(q,p%q);}
ll qmul(ll p,ll q,ll mo){ll f=0;while(q){if(q&1)f=(f+p)%mo;p=(p+p)%mo;q>>=1;}return f;}
ll qpow(ll p,ll q){ll f=1;while(q){if(q&1)f=f*p;p=p*p;q>>=1;}return f;}
int n,m,k,t;
struct samnode{
    samnode *son[10] , *f;
    int l ;
}*root,*last,sam[maxn*40];
int cnt;
void init(){
    root = last = &sam[cnt=0];
}
void add(int x,samnode *last)
{
    samnode *p = &sam[++cnt] , *jp=last;
    p->l = jp->l+1;
    last = p;
    for( ; jp&&!jp->son[x] ; jp=jp->f) jp->son[x]=p;
    if(!jp) p->f = root;
    else{
        if(jp->l+1 == jp->son[x]->l) p->f = jp->son[x];
        else{
            samnode *r = &sam[++cnt] , *q = jp->son[x];
            *r = *q;
            r->l = jp->l+1;
            q->f = p->f = r;
            for( ; jp && jp->son[x]==q ; jp=jp->f) jp->son[x]=r;
        }
    }
}
vi e[maxn];
int du[maxn],c[maxn];
void dfs(int x,int y,samnode *rt)
{
    int i;
    add(c[x],rt);
    rt=rt->son[c[x]];
    rep(i,0,e[x].size()-1)
    {
        int z=e[x][i];
        if(z==y)continue;
        dfs(z,x,rt);
    }
}
int main(){
    int i,j;
    init();
    scanf("%d%d",&n,&m);
    rep(i,1,n)scanf("%d",&c[i]);
    rep(i,1,n-1)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        du[x]++,du[y]++;
        e[x].pb(y),e[y].pb(x);
    }
    rep(i,1,n)if(du[i]==1)dfs(i,0,root);
    ll ret=0;
    rep(i,1,cnt)ret+=sam[i].l-sam[i].f->l;
    printf("%lld\n",ret);
    return 0;
}