[TJOI2015]弦论 - 后缀自动机

Posted 2021-03-13 mollnn

下了狠心开始做SAM的题目了……

（中间因为傻逼26分写错被卡，进来的时候记得把自己的 cnt 减掉）

// TJOI2015 XIAN LUN
#include <bits/stdc++.h>
using namespace std;
const int Maxn = 2000005;
struct Suffix_Automata {
    int maxlen[Maxn], trans[Maxn][26], link[Maxn], Size, Last;
    int t[Maxn], a[Maxn], cnt[Maxn], f[Maxn];
    Suffix_Automata() { Size = Last = 1; }
    inline void Extend(int id) {
        int cur = (++ Size), p;
        maxlen[cur] = maxlen[Last] + 1;
        cnt[cur] = 1;
        for (p = Last; p && !trans[p][id]; p = link[p]) trans[p][id] = cur;
        if (!p) link[cur] = 1;
        else {
            int q = trans[p][id];
            if (maxlen[q] == maxlen[p] + 1) link[cur] = q;
            else {
                int clone = (++ Size);
                maxlen[clone] = maxlen[p] + 1;
                for(int i=0;i<26;i++) trans[clone][i] = trans[q][i];
                link[clone] = link[q];
                for (; p && trans[p][id] == q; p = link[p]) trans[p][id] = clone;
                link[cur] = link[q] = clone;
            }
        }
        Last = cur;
    }
    void CalcEndposSize() {
        memset(t, 0, sizeof t);
        for(int i=1; i<=Size; i++) t[maxlen[i]]++;
        for(int i=1; i<=Size; i++) t[i]+=t[i-1];
        for(int i=1; i<=Size; i++) a[t[maxlen[i]]--]=i;
        for(int i=Size; i>=1; --i) cnt[link[a[i]]]+=cnt[a[i]];
        cnt[1] = 0;
    }
    void DFS(int p) {
        for(int i=0;i<26;i++) {
            if(trans[p][i]) {
                if(f[trans[p][i]]==0) DFS(trans[p][i]);
                f[p]+=f[trans[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(trans[p][i]) {
                if(f[trans[p][i]]>=k) {
                    cout<<(char)(i+'a');
                    Go(trans[p][i],k);
                    return;
                }
                else {
                    k-=f[trans[p][i]];
                }
            }
        }
        if(p==1) cout<<-1;
    }
    void Calc(int k) {
        DFS(1);
        Go(1,k);
    }
} sam;

int main() {
    string str;
    cin>>str;
    int t,k;
    cin>>t>>k;
    for(int i=0;i<str.length();i++)
        sam.Extend(str[i]-'a');
    sam.CalcEndposSize();
    if(t==0) {
        for(int i=2; i<=sam.Size; i++)
            sam.cnt[i] = 1;
    }
    sam.Calc(k);
}