CF123D String (SAM)

CF123D String

Mean

给出字符串 \(s\)，定义子串 \(a\) 在 \(s\) 中的出现次数为 \(cnt(a)\)，求\(\sum \frac{cnt(a)(cnt(a)+1)}{2}\)。

Sol

SAM裸题。

建出SAM后答案为\(\sum(node[i].len-node[node[i].fa].len)*(siz[i]*(siz[i]+1))/2\)

Code

#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;
typedef long long ll;
const int N = 3e5+10;
int tot = 1, last = 1;
struct Node
{
    int len, fa;
    int ch[26];
}node[N];
char str[N];
ll f[N], ans;
int h[N], e[N], ne[N], idx;

void extend(int c)
{
    int p = last, np = last = ++ tot;
    f[tot] = 1;
    node[np].len = node[p].len + 1;
    for (; p && !node[p].ch[c]; p = node[p].fa) node[p].ch[c] = np;
    if (!p) node[np].fa = 1;
    else
    {
        int q = node[p].ch[c];
        if (node[q].len == node[p].len + 1) node[np].fa = q;
        else
        {
            int nq = ++ tot;
            node[nq] = node[q], node[nq].len = node[p].len + 1;
            node[q].fa = node[np].fa = nq;
            for (; p && node[p].ch[c] == q; p = node[p].fa) node[p].ch[c] = nq;
        }
    }
}

void add(int a, int b)
{
    e[idx] = b, ne[idx] = h[a], h[a] = idx ++ ;
}

void dfs(int u)
{
    for (int i = h[u]; ~i; i = ne[i])
    {
        dfs(e[i]);
        f[u] += f[e[i]];
    }
}

int main()
{
    scanf("%s", str);
    for (int i = 0; str[i]; i ++ ) extend(str[i] - 'a');
    memset(h, -1, sizeof h);
    for (int i = 2; i <= tot; i ++ ) add(node[i].fa, i);
    dfs(1);
    for(int i=2;i<=tot;++i){
        ans = ans + (node[i].len-node[node[i].fa].len)*f[i]*(f[i]+1)/2;
    }
    printf("%lld\n", ans);

    return 0;
}