「TJOI2015」弦论 后缀自动机求第k小串

「TJOI2015」弦论

题意

对于一个给定长度为\(N\)的字符串，求它的第\(K\)小子串，若\(T\)为\(0\)表示不同位置的相同子串算作一个，否则算作多个，子串数目不足\(K\)个，则输出\(-1\)。

\(n\le 5\cdot 10^5\)

分析

建sam，按拓扑序动态规划求出每个结点后有多少个子串，然后贪心求解即可。

Code

#include<bits/stdc++.h>
#define rep(i,x,n) for(int i=x;i<=n;i++)
#define per(i,n,x) for(int i=n;i>=x;i--)
#define sz(a) int(a.size())
#define rson mid+1,r,p<<1|1
#define pii pair<int,int>
#define lson l,mid,p<<1
#define ll long long
#define pb push_back
#define mp make_pair
#define se second
#define fi first
using namespace std;
const double eps=1e-8;
const int mod=1e9+7;
const int N=2e6+10;
const int inf=1e9;
int n,T,k;
char s[N];
struct SAM{
    int last,cnt;int ch[N][26],fa[N],len[N],sum[N],id[N],sz[N];
    ll g[N];
    int newnode(){
        ++cnt;
        return cnt;
    }
    void insert(int c){
        int p=last,np=newnode();last=np;len[np]=len[p]+1;
        for(;p&&!ch[p][c];p=fa[p]) ch[p][c]=np;
        if(!p) fa[np]=1;
        else {
            int q=ch[p][c];
            if(len[q]==len[p]+1) fa[np]=q;
            else{
                int nq=newnode();
                len[nq]=len[p]+1;
                memcpy(ch[nq],ch[q],sizeof ch[q]);
                fa[nq]=fa[q],fa[q]=fa[np]=nq;
                for(;ch[p][c]==q;p=fa[p]) ch[p][c]=nq;
            }
        }
        sz[np]=1;
    }
    void init(){
        last=cnt=1;
    }
    void gao(){
        rep(i,1,cnt) sum[len[i]]++;
        rep(i,1,cnt) sum[i]+=sum[i-1];
        rep(i,1,cnt) id[sum[len[i]]--]=i;
        per(i,cnt,1) sz[fa[id[i]]]+=sz[id[i]];
        rep(i,1,cnt) g[i]=T?sz[i]:1;
        g[1]=0;
        per(i,cnt,1){
            int u=id[i];
            for(int j=0;j<26;j++) if(ch[u][j]){
                int x=ch[u][j];
                g[u]+=g[x];
            }
        }
    }
    void solve(int k){
        if(g[1]<k){
            puts("-1");
            return;
        }
        int u=1;
        while(k>0){
            for(int i=0;i<26;i++) if(ch[u][i]){
                int x=ch[u][i];
                if(g[x]<k) k-=g[x];
                else{
                    putchar(i+'a');
                    k-=T?sz[x]:1;
                    u=x;
                    break;
                }
            }
        }
        puts("");
    }
}sam;
int main(){
    scanf("%s",s+1);
    n=strlen(s+1);
    sam.init();
    rep(i,1,n) sam.insert(s[i]-'a');
    scanf("%d%d",&T,&k);
    sam.gao();
    sam.solve(k);
    return 0;
}