BZOJ3926 (後綴自動機)

BZOJ3926 諸神眷顧的幻想鄉

Problem :
給一個n個節點的樹（n<=10^5), 每一個點有一種顏色(c<=10), 詢問全部點對之間路徑組成字符串的種類。保證葉子節點小於等於20.
Solution :
分別從每一個葉子節點開始遍歷整棵樹，將遍歷到的字符串加入後綴自動機。
只須要修改後綴自動機的last，就能夠實現添加多串。

#include <bits/stdc++.h>
using namespace std;

const int N = 1e5 * 15 * 2;
vector <int> eg[N];
int cl[N];
int deg[N];
int n, c;

struct Suffix_Automanon
{
    int nt[N][15], fail[N], a[N];
    int p, q, np, nq;
    int tot, root;
    int newnode(int len)
    {
        for (int i = 0; i < 15; ++i) nt[tot][i] = -1;
        fail[tot] = -1; a[tot] = len;
        return tot++;
    }
    void clear()
    {
        tot = 0;
        root = newnode(0);
    }
    void insert(int ch, int &last)
    {
        p = last; last = np = newnode(a[p] + 1); 
        for (; ~p && nt[p][ch] == -1; p = fail[p]) nt[p][ch] = np;
        if (p == -1) fail[np] = root;
        else
        {
            q = nt[p][ch];
            if (a[q] == a[p] + 1) fail[np] = q;
            else
            {
                nq = newnode(a[p] + 1);
                for (int i = 0; i < 15; ++i) nt[nq][i] = nt[q][i];
                fail[nq] = fail[q]; fail[q] = fail[np] = nq;
                for (; ~q && nt[p][ch] == q; p = fail[p]) nt[p][ch] = nq;
            }
        }
    }
    void solve()
    {
        long long ans = 0;
        for (int i = 1; i < tot; ++i)
            ans += a[i] - a[fail[i]];
        cout << ans << "\n";
    }
}sam;
void dfs(int u, int fa, int now)
{
    sam.insert(cl[u], now);
    for (int i = 0; i < (int)eg[u].size(); ++i)
    {
        int v = eg[u][i];
        if (v != fa) dfs(v, u, now);
    }
}
int main()
{
    cin.sync_with_stdio(0);
    cin >> n >> c;  
    for (int i = 1; i <= n; ++i) cin >> cl[i];
    for (int i = 1; i <= n; ++i) eg[i].clear(), deg[i] = 0;
    for (int i = 1; i <  n; ++i)
    {
        int u, v; 
        cin >> u >> v;
        eg[u].push_back(v);
        eg[v].push_back(u);
        deg[u]++; deg[v]++;
    }
    sam.clear();
    for (int i = 1; i <= n; ++i)
        if (deg[i] == 1)
            dfs(i, 0, 0);
    sam.solve();
}