SP7258 SUBLEX - Lexicographical Substring Search

Description

Little Daniel loves to play with strings! He always finds different ways to have fun with strings! Knowing that, his friend Kinan decided to test his skills so he gave him a string S and asked him Q questions of the form:

If all distinct substrings of string S were sorted lexicographically, which one will be the K-th smallest?

After knowing the huge number of questions Kinan will ask, Daniel figured out that he can't do this alone. Daniel, of course, knows your exceptional programming skills, so he asked you to write him a program which given S will answer Kinan's questions.
Example:

S = "aaa" (without quotes)
substrings of S are "a" , "a" , "a" , "aa" , "aa" , "aaa". The sorted list of substrings will be:
"a", "aa", "aaa".

Solution

在后缀自动机的DAG上DP，贪心地选取节点

#include<iostream>
#include<cstring>
#include<cstdio>
#include<queue>
using namespace std;
int n,las=1,tot=1,T,tot2,head[1800010],du[1800010],dp[1800010];
char s[900005];
struct SAM
{
    int ch[26],fa,len;
}sam[5000010];
struct Edge
{
    int to,nxt;
}edge[1800010];
queue<int>q;
inline int read()
{
    int w=0,f=1;
    char ch=0;
    while(ch<'0'||ch>'9'){if(ch=='-') f=-1; ch=getchar();}
    while(ch>='0'&&ch<='9')w=(w<<1)+(w<<3)+ch-'0',ch=getchar();
    return w*f;
}
void insert(int c)
{
    int p=las,np=las=++tot;
    sam[np].len=sam[p].len+1;
    for(;p&&!sam[p].ch[c];p=sam[p].fa) sam[p].ch[c]=np;
    if(!p) sam[np].fa=1;
    else
    {
        int q=sam[p].ch[c];
        if(sam[q].len==sam[p].len+1) sam[np].fa=q;
        else
        {
            int nq=++tot;
            sam[nq]=sam[q],sam[nq].len=sam[p].len+1,sam[q].fa=sam[np].fa=nq;
            for(;p&&sam[p].ch[c]==q;p=sam[p].fa) sam[p].ch[c]=nq;
        }
    }
}
int main()
{
    scanf("%s",s+1),n=strlen(s+1),T=read();
    for(int i=1;i<=n;i++) insert(s[i]-'a');
    for(int i=1;i<=tot;i++) for(int j=0;j<26;j++) if(sam[i].ch[j]) edge[++tot2]=(Edge){i,head[sam[i].ch[j]]},head[sam[i].ch[j]]=tot2,++du[i];
    for(int i=1;i<=tot;i++) if(!du[i]) q.push(i);
    while(q.size())
    {
        int u=q.front();
        q.pop(),dp[u]=1;
        for(int i=0;i<26;i++) if(sam[u].ch[i]) dp[u]+=dp[sam[u].ch[i]];
        for(int i=head[u];i;i=edge[i].nxt)
        {
            int v=edge[i].to;
            du[v]--;
            if(!du[v]) q.push(v);
        }
    }
    for(;T;T--)
    {
        int k=read(),p=1;
        while(k) for(int i=0;i<26;i++) if(sam[p].ch[i])
            if(dp[sam[p].ch[i]]>=k)
            {
                putchar(i+'a'),p=sam[p].ch[i],--k;break;
            }
            else k-=dp[sam[p].ch[i]];
        putchar(10);
    }
    return 0;
}