[ZJOI2015] 诸神眷顾的幻想乡

Description

给定一棵叶子结点数量 (le 20) 的树，在树上任选一条路，问有多少种颜色序列不同的选择。

Solution

一条路径一定属于以某个叶子为根的树上所有的直链的集合

于是我们把以每个叶子为根的树 DFS 一遍，当作一个 Trie 插入广义后缀自动机即可

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int N = 2000005;

int n,c,a[N],t1,t2,t3,vis[N],pos[N],d[N];
vector <int> g[N];

struct SAM {
    int len[N], ch[N][10], fa[N], ind, last;
    SAM() { ind = last = 1; }
    inline int extend(int id) {
        if(ch[last][id] && len[last]+1==len[ch[last][id]]) return ch[last][id]; //!
        int cur = (++ ind), p, tmp, flag = 0; //!
        len[cur] = len[last] + 1;
        for (p = last; p && !ch[p][id]; p = fa[p]) ch[p][id] = cur;
        if (!p) fa[cur] = 1;
        else {
            int q = ch[p][id];
            if (len[q] == len[p] + 1) fa[cur] = q;
            else {
                if(p==last) flag=1; //!
                tmp = (++ ind);
                len[tmp] = len[p] + 1;
                for(int i=0;i<10;i++) ch[tmp][i] = ch[q][i];
                fa[tmp] = fa[q];
                for (; p && ch[p][id] == q; p = fa[p]) ch[p][id] = tmp;
                fa[cur] = fa[q] = tmp;
            }
        }
        last = cur;
        return flag ? tmp : cur;//!
    }
    void dfs(int p)
    {
        vis[p]=1;
        for(int q:g[p])
        {
            if(!vis[q])
            {
                last=pos[p];
                pos[q]=extend(a[q]);
                dfs(q);
            }
        }
    }
    void ins(int rt)
    {
        memset(vis,0,sizeof vis);
        memset(pos,0,sizeof pos);
        last=1;
        pos[rt]=extend(a[rt]);
        dfs(rt);
    }
    int getans()
    {
        int ans=0;
        for(int i=1;i<=ind;i++) ans+=len[i]-len[fa[i]];
        return ans;
    }
} sam;

signed main() {
    ios::sync_with_stdio(false);
    cin>>n>>c;
    for(int i=1;i<=n;i++) cin>>a[i];
    for(int i=1;i<n;i++)
    {
        cin>>t1>>t2;
        g[t1].push_back(t2);
        g[t2].push_back(t1);
        d[t1]++;
        d[t2]++;
    }
    for(int i=1;i<=n;i++)
    {
        if(d[i]==1)
        {
            sam.ins(i);
        }
    }
    cout<<sam.getans()<<endl;
}