BZOJ 3926 && ZJOI 2015 诸神眷顾的幻想乡 (广义后缀自动机)

3926: [Zjoi2015]诸神眷顾的幻想乡

Time Limit: 10 Sec  Memory Limit: 512 MB

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<=n<=100000, 1<=c<=10。
对于15%的数据，n<=2000。
另有5%的数据，所有空地都至多与两个空地相邻。
另有5%的数据，除一块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻。
另有5%的数据，除某两块空地与三个空地相邻外，其他空地都分别至多与两个空地相邻

算法讨论：

由于叶子节点不超过20个，所以我们分别从每个叶子出发建立广义后缀自动机。注意树上的自动建立要时刻清楚last节点是哪个。

然后就可以发现，树的任何一个子串，都是自动机上的某个节点到其子孙的路径。然后本质不同的子串数量就是每个点和其pre指针的len差值。

加和就可以啦~~~

代码：

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

using namespace std;
const int N = 100000 + 5;
const int C = 10;

struct State {
  int len, pre, next[C];
}st[N * 40];

struct SuffixAutomaton {
  int sz;
  void Init() {
    sz = 1;
    st[sz].pre = -1; st[sz].len = 0;
    sz ++;
  }
  int add(int c, int last) {
    int cur = sz ++, p;
    st[cur].len = st[last].len + 1;
    for(p = last; p != -1 && !st[p].next[c]; p = st[p].pre)
      st[p].next[c] = cur;
    if(p == -1) st[cur].pre = 1;
    else {
      int q = st[p].next[c];
      if(st[q].len == st[p].len + 1) st[cur].pre = q;
      else {
        int cle = sz ++;
        st[cle].pre = st[q].pre;
        st[cle].len = st[p].len + 1;
        for(int i = 0; i < C; ++ i) st[cle].next[i] = st[q].next[i];
        for(; p != -1 && st[p].next[c] == q; p = st[p].pre)
          st[p].next[c] = cle;
        st[q].pre = st[cur].pre = cle;
      }
    }
    last = cur;
    return last;
  }
  void solve() {
    long long ans = 0;
    for(int i = 2; i < sz; ++ i)
        ans = (long long) ans + st[i].len - st[st[i].pre].len;
    printf("%lld\n", ans);
  }
}sam;

int n, c, cnt;
int head[N << 1], w[N << 1], du[N];
struct Edge {
  int from, to, next;
}edges[N << 1];

void insert(int from, int to) {
  ++ cnt;
  edges[cnt].from = from; edges[cnt].to = to;
  edges[cnt].next = head[from]; head[from] = cnt;
}

void dfs(int u, int f, int last) {
  last = sam.add(w[u], last);
  for(int i = head[u]; i; i = edges[i].next) {
    int v = edges[i].to;
    if(v != f) dfs(v, u, last);
  }
}

int main() {
  //freopen("zjoi15_substring.in", "r", stdin);
  //freopen("zjoi15_substring.out", "w", stdout);

  int x, y;
  scanf("%d%d", &n, &c);
  for(int i = 1; i <= n; ++ i) scanf("%d", &w[i]);
  for(int i = 1; i < n; ++ i) {
    scanf("%d%d", &x, &y);
    insert(x, y); insert(y, x);
    du[x] ++; du[y] ++;
  }
  sam.Init();
  for(int i = 1; i <= n; ++ i)
    if(du[i] == 1) { dfs(i, 0, 1); }
  sam.solve();

  //fclose(stdin); fclose(stdout);
  return 0;
}