【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\le n\le 100000, 1\le c\le 10\)。

對於\(15\%\)的資料，\(n\le 2000\)。

另有\(5\%\)的資料，所有空地都至多與兩個空地相鄰。

另有\(5\%\)的資料，除一塊空地與三個空地相鄰外，其他空地都分別至多與兩個空地相鄰。

另有\(5\%\)的資料，除某兩塊空地與三個空地相鄰外，其他空地都分別至多與兩個空地相鄰

壓根沒搞懂什麼是對的廣義SAM，什麼是錯的廣義SAM，我的程式碼說不定是假的。

具體思路，直接對每個度為1的點建廣義SAM，就是每次處理的las是原樹父親節點建出的那個點要處理，其他和SAM一樣。直接把所有點的廣義SAM建到一起就行了。

Code:

#include <cstdio>
#include <cstring>
#define ll long long
const int N=1e5+10;
const int M=4e6+10;
int ch[M][10],len[M],par[M],tot=1;
int in[N],s[N];
int extend(int p,int c)
{
    int now=++tot;
    len[now]=len[p]+1;
    while(p&&!ch[p][c]) ch[p][c]=now,p=par[p];
    if(!p) par[now]=1;
    else
    {
        int x=ch[p][c];
        if(len[x]==len[p]+1) par[now]=x;
        else
        {
            int y=++tot;
            len[y]=len[p]+1,par[y]=par[x];
            memcpy(ch[y],ch[x],sizeof ch[x]);
            while(p&&ch[p][c]==x) ch[p][c]=y,p=par[p];
            par[now]=par[x]=y;
        }
    }
    return now;
}
int head[N],to[N<<1],Next[N<<1],cnt;
void add(int u,int v)
{
    to[++cnt]=v,Next[cnt]=head[u],head[u]=cnt;
}
void dfs(int now,int fa,int las)
{
    int bee=extend(las,s[now]);
    for(int v,i=head[now];i;i=Next[i])
        if((v=to[i])^fa)
            dfs(v,now,bee);
}
int main()
{
    int n,c;scanf("%d%d",&n,&c);
    for(int i=1;i<=n;i++) scanf("%d",s+i);
    for(int u,v,i=1;i<n;i++) scanf("%d%d",&u,&v),++in[u],++in[v],add(u,v),add(v,u);
    for(int i=1;i<=n;i++) if(in[i]==1) dfs(i,0,1);
    ll ans=0;
    for(int i=1;i<=tot;i++) ans=ans+len[i]-len[par[i]];
    printf("%lld\n",ans);
    return 0;
}

2019.1.10