hihoCoder1445—后缀自动机二·重复旋律5

题目链接

时间限制: 10000ms

单点时限: 2000ms

内存限制: 512MB

描述

小Hi平时的一大兴趣爱好就是演奏钢琴。我们知道一个音乐旋律被表示为一段数构成的数列。

现在小Hi想知道一部作品中出现了多少不同的旋律？

解题方法提示

输入

共一行，包含一个由小写字母构成的字符串。字符串长度不超过 1000000。

输出

一行一个整数，表示答案。

样例输入

aab

样例输出

5

解题思路：

不同的旋律就是字符串中不同的子串个数，在这里对S建立后缀自动机，自动机上从初始状态出发的路径数就是不同的子串个数。因为后缀自动机是一个有向无环图，这里可以用dp记录从某个状态出发的路径数，于是有dp[v] = sum(dp[u]+1) v-->u。

这里我使用的是记忆化搜索。

代码：

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstring>
#include <string>
#include <map>

using namespace std;

typedef long long llt;

const int N = 1000010;

struct state{
    int len,link;  //link为后缀链接，len表示结点到起始结点的最大长度
    map<char,int>next;  //st[0].next[c] = 1表示0结点经过状态c转移到节点1
}st[N*2];

llt dp[N*2];

int sz,last;

void sa_init()
{
    sz = last = 0;
    st[0].len = 0;
    st[0].link = -1;
    ++sz;

    memset(dp,0,sizeof(dp));

    for(int i = 0; i < N; ++i)
        st[i].next.clear();
}

void sa_extend(char c)
{
    int cur = sz++;         //新添加的状态结点
    st[cur].len = st[last].len+1;
    int p;
    for(p = last; p != -1 && !st[p].next.count(c); p = st[p].link)
        st[p].next[c] = cur;
    if(p == -1)
        st[cur].link = 0;
    else{
        int q = st[p].next[c];
        if(st[p].len+1 == st[q].len)
            st[cur].link = q;
        else{
            int clone = sz++;
            st[clone].len = st[p].len+1;
            st[clone].next = st[q].next;
            st[clone].link = st[q].link;
            for(; p != -1 && st[p].next[c] == q; p = st[p].link)
                st[p].next[c] = clone;
            st[q].link = st[cur].link = clone;
        }
    }
    last = cur;
}

llt dfs(int s)
{
    llt ans = 0;

    if(dp[s]) return dp[s];

    for(auto it = st[s].next.begin(); it != st[s].next.end(); ++it){
        ans += dfs((*it).second)+1;
    }

    dp[s] = ans;

    return ans;
}

int main()
{
    sa_init();
    string str;
    cin >> str;
    for(int i = 0; i < str.length(); ++i)
        sa_extend(str[i]);

    printf("%lld\n",dfs(0));
    return 0;
}