[CF128B] String

[CF128B] String - SAM,dp

Description

求给定字符串的字典序第 (k) 小子串（允许重复）。

Solution

SAM 上 dp，预处理出 (endpos) 集合大小 (cnt[p])，则 (f[p]=sum_{p o q} f[q] + cnt[p])

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 1000005;
struct SAM
{
    int len[N], ch[N][26], fa[N], ind, last;
    int t[N], a[N], cnt[N], f[N];
    SAM() { ind = last = 1; }
    inline void extend(int id)
    {
        int cur = (++ind), p;
        len[cur] = len[last] + 1;
        cnt[cur] = 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
            {
                int tmp = (++ind);
                len[tmp] = len[p] + 1;
                for (int i = 0; i < 26; 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;
    }
    void calcEndpos()
    {
        memset(t, 0, sizeof t);
        for (int i = 1; i <= ind; i++)
            t[len[i]]++;
        for (int i = 1; i <= ind; i++)
            t[i] += t[i - 1];
        for (int i = 1; i <= ind; i++)
            a[t[len[i]]--] = i;
        for (int i = ind; i >= 1; --i)
            cnt[fa[a[i]]] += cnt[a[i]];
        cnt[1] = 0;
    }
    void DFS(int p)
    {
        for (int i = 0; i < 26; i++)
        {
            if (ch[p][i])
            {
                if (f[ch[p][i]] == 0)
                    DFS(ch[p][i]);
                f[p] += f[ch[p][i]];
            }
        }
        f[p] += cnt[p];
    }
    void Go(int p, int k)
    {
        k -= cnt[p];
        for (int i = 0; i < 26 && k > 0; i++)
        {
            if (ch[p][i])
            {
                if (f[ch[p][i]] >= k)
                {
                    cout << (char)(i + 'a');
                    Go(ch[p][i], k);
                    return;
                }
                else
                {
                    k -= f[ch[p][i]];
                }
            }
        }
        if (k > 0)
        {
            cout << "No such line.";
            exit(0);
        }
    }
} sam;

signed main()
{
    ios::sync_with_stdio(false);
    string str;
    cin >> str;
    int t, k;
    for (int i = 0; i < str.length(); i++)
        sam.extend(str[i] - 'a');
    sam.calcEndpos();
    sam.DFS(1);
    cin >> k;
    sam.Go(1, k);
}