【ZJOI2015】諸神眷顧的幻想鄉

Description

Input

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

Data Constraint

n<=100000,c<=10

Solution

發現葉子只有10個，那麼可以以每個葉子爲根，建trie，然後建廣義後綴自動機
對於每個狀態，代表的字符串個數就是max-min+1

Code

#include<cstdio>
#include<cstring>
#include<algorithm>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
#define N 101000
#define ll long long
using namespace std;
int n,str,s[N],d[N*20][2],tot=0,last[N],next[N*2],to[N*2],a[N][21],rd[N];
ll ans=0;
struct SAM{
    int len,fail,size;
    int to[10];
}t[N*50];
void putin(int x,int y)
{
    next[++tot]=last[x];last[x]=tot;to[tot]=y;
}
void add(int x,int cm,int nu,int last)
{
    int p=last,np=++tot;
    t[np].len=t[p].len+1;t[np].size=1;a[cm][nu]=np;
    for(;p&&t[p].to[x]==0;p=t[p].fail) t[p].to[x]=np;
    if(p==0) t[np].fail=1;
    else
    {
        int q=t[p].to[x];
        if(t[p].len+1==t[q].len) t[np].fail=q;
        else
        {
            int nq=++tot;t[nq]=t[q];
            t[nq].len=t[p].len+1;t[q].fail=t[np].fail=nq;
            for(;p&&t[p].to[x]==q;p=t[p].fail) t[p].to[x]=nq;
        }
    }
}
void pre()
{
    int he=0,ta=0,num=0;
    fo(i,1,n) if(rd[i]==1) add(s[i],i,++num,1),d[++ta][0]=i,d[ta][1]=num;
    while(he<ta)
    {
        int x=d[++he][0],nu=d[he][1];
        for(int i=last[x];i;i=next[i])
        {
            int y=to[i];if(a[y][nu]) continue;
            add(s[y],y,nu,a[x][nu]);d[++ta][0]=y;d[ta][1]=nu;
        }
    }
}
int main()
{
    scanf("%d%d",&n,&str);
    fo(i,1,n) scanf("%d",&s[i]);
    fo(i,1,n-1)
    {
        int x,y;scanf("%d%d",&x,&y);
        putin(x,y);putin(y,x);rd[x]++;rd[y]++;
    }
    tot=1;pre();
    fo(i,1,tot) ans+=(t[i].len-t[t[i].fail].len);
    printf("%lld\n",ans);
}