BZOJ 3998 后缀数组

//By SiriusRen
#include <bits/stdc++.h>
using namespace std;
const int N=1000050;
int n,cntA[N],cntB[N],A[N],B[N],rk[N],sa[N],tsa[N],ht[N],T,K;
long long sum[N],sum1[N];
char s[N];
void SA(){
    for(int i=1;i<=n;i++)cntA[s[i]]++;
    for(int i=1;i<=256;i++)cntA[i]+=cntA[i-1];
    for(int i=n;i;i--)sa[cntA[s[i]]--]=i;
    rk[sa[1]]=1;
    for(int i=2;i<=n;i++)rk[sa[i]]=rk[sa[i-1]]+(s[sa[i]]!=s[sa[i-1]]);
    for(int l=1;rk[sa[n]]<n;l<<=1){
        memset(cntA,0,sizeof(cntA));
        memset(cntB,0,sizeof(cntB));
        for(int i=1;i<=n;i++)cntA[A[i]=rk[i]]++,cntB[B[i]=(i+l<=n?rk[i+l]:0)]++;
        for(int i=1;i<=n;i++)cntA[i]+=cntA[i-1],cntB[i]+=cntB[i-1];
        for(int i=n;i;i--)tsa[cntB[B[i]]--]=i;
        for(int i=n;i;i--)sa[cntA[A[tsa[i]]]--]=tsa[i];
        rk[sa[1]]=1;
        for(int i=2;i<=n;i++)rk[sa[i]]=rk[sa[i-1]]+(A[sa[i]]!=A[sa[i-1]]||B[sa[i]]!=B[sa[i-1]]);
    }
    for(int i=1,j=0;i<=n;i++){
        j=j?j-1:0;
        while(s[i+j]==s[sa[rk[i]-1]+j])j++;
        ht[rk[i]]=j;
    }
}
void print(int l,int r){for(int i=l;i<=r;i++)putchar(s[i]);}
bool check(int p){
    int tempans=1,l=1,r=n;
    while(l<=r){
        int mid=(l+r)>>1;
        if(sum[mid]>=p)r=mid-1;
        else tempans=mid+1,l=mid+1;
    }
    int hi=p-sum[tempans-1]+ht[tempans],tot=hi+sum1[tempans-1];
    if(tot>=K)return 1;
    for(int i=tempans+1;i<=n;i++){
        hi=min(hi,ht[i]);
        if(!hi)break;
        tot+=hi;
        if(tot>=K)return 1;
    }return 0;
}
void solve(){
    for(int i=1,t;i<=n;i++,K-=t){
        t=n-sa[i]-ht[i]+1;
        if(K<=t){print(sa[i],sa[i]+K-1+ht[i]);return;}
    }puts("-1");
}
signed main(){
    scanf("%s%d%d",s+1,&T,&K),n=strlen(s+1),SA();
    if(!T)solve();
    else{
        for(int i=1;i<=n;i++)sum[i]=sum[i-1]+n-ht[i]-sa[i]+1,sum1[i]=sum1[i-1]+n-sa[i]+1;
        if(sum1[n]<K){puts("-1");return 0;}
        int l=1,r=K+1,ans;
        while(l<=r){
            int mid=(l+r)>>1;
            if(check(mid))ans=mid,r=mid-1;
            else l=mid+1;
        }K=ans;solve();
    }
}