Description:

给定一颗含有 个节点的树,保证叶子节点个数小于等于 ,每个节点有一个颜色在 之间,定义颜色序列 表示在树上从节点 的路径中依次经过的点的颜色形成的序列,问一共有多少种颜色序列存在于树上的颜色序列中。

Solution:

我不得不说看原题很容易以为是点的度数小于 ,然后就。。。。

既然叶子节点只有 那么可以对于每个叶子节点作根构建 树,然后把所有 树合起来(事实上只需要在同一棵 树上不断插入就行了,并不需要真正的合并),之后对大的 树构建后缀自动机,然后拓扑排序求出这个后缀自动机上本质不同的子串有多少个,经典问题。

Code:

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

using namespace std;

int N,C;
int Color[100010]={0};
struct bian_
{
    int to;
    int next;
}bian[200010]={{0,0}};
int First[100010]={0};
int du[100010]={0};
struct Trie
{
    int son[10];
    int fa;
}trie[2000010]={{{0},0}};
int tp=0;

struct SAM
{
    int son[10];
    int fail;
    int lenth;
    long long G;
}pot[4000010]={{{0},0,0,0}};
int pp=0;
int New[2000010]={0};

int dui[2000010]={0};
int duip=0;

int S[4000010]={0};
int EN[100010]={0};

void dfs(int cnt1,int fa,int cnt2)
{
    for(int i=First[cnt1];i!=0;i=bian[i].next)
    {
        int u=bian[i].to;
        if(u==fa) continue;
        if(trie[cnt2].son[Color[u]]==0)
        {
            trie[cnt2].son[Color[u]]=++tp;
            trie[tp].fa=cnt2;
        }
        dfs(u,cnt1,trie[cnt2].son[Color[u]]);
    }
    return;
}

void Add(int p,int q,int k)
{
    bian[k].to=q;
    bian[k].next=First[p];
    First[p]=k;
    return;
}

void Build_sam(int t,int c)
{
    int cnt=New[trie[t].fa];New[t]=++pp;
    pot[New[t]].lenth=pot[cnt].lenth+1;
    for(;cnt!=0 && pot[cnt].son[c]==0;cnt=pot[cnt].fail)
        pot[cnt].son[c]=New[t];
    if(cnt==0)
        pot[New[t]].fail=1;
    else
    {
        int q=pot[cnt].son[c];
        if(pot[cnt].lenth+1==pot[q].lenth)
            pot[New[t]].fail=q;
        else
        {
            pp++;pot[pp]=pot[q];pot[pp].lenth=pot[cnt].lenth+1;
            pot[New[t]].fail=pot[q].fail=pp;
            for(;cnt!=0 && pot[cnt].son[c]==q;cnt=pot[cnt].fail)
                pot[cnt].son[c]=pp;
        }
    }
    return;
}

void bfs_and_makesam()
{
    dui[++duip]=1;
    New[1]=++pp;
    for(int i=1;i<=duip;i++)
    {
        int u=dui[i];
        for(int j=0;j<C;j++)
        {
            if(trie[u].son[j]!=0)
            {
                dui[++duip]=trie[u].son[j];
                Build_sam(trie[u].son[j],j);
            }
        }
    }
    return;
}

int main()
{
    cin>>N>>C;
    for(int i=1;i<=N;i++)
        scanf("%d",&Color[i]);
    for(int i=1;i<N;i++)
    {
        int p,q;
        scanf("%d%d",&p,&q);
        Add(p,q,(i<<1)-1);
        Add(q,p,i<<1);
        du[p]++,du[q]++;
    }
    tp++;
    for(int i=1;i<=N;i++)
        if(du[i]==1)
        {
            if(trie[1].son[Color[i]]==0)
            {
                trie[1].son[Color[i]]=++tp;
                trie[tp].fa=1;
            }
            dfs(i,0,trie[1].son[Color[i]]);
        }
    bfs_and_makesam();
    for(int i=1;i<=pp;i++)
        EN[pot[i].lenth]++;
    for(int i=0;i<=N;i++)
        EN[i]+=EN[i-1];
    for(int i=1;i<=pp;i++)
        S[EN[pot[i].lenth]--]=i;
    for(int i=pp;i>=1;i--)
    {
        int u=S[i];
        for(int j=0;j<C;j++)
        {
            if(pot[u].son[j]!=0)
                pot[u].G+=pot[pot[u].son[j]].G;
        }
        pot[u].G++;
    }
    cout<<pot[1].G-1<<endl;
    return 0;
}