BZOJ 3998 [TJOI2015]弦论

Description

对于一个给定长度为N的字符串，求它的第K小子串是什么。

Input

第一行是一个仅由小写英文字母构成的字符串S
第二行为两个整数T和K，T为0则表示不同位置的相同子串算作一个。T=1则表示不同位置的相同子串算作多个。K的意义如题所述。

Output

输出仅一行，为一个数字串，为第K小的子串。如果子串数目不足K个，则输出-1

Sample Input

aabc
0 3

Sample Output

aab

HINT

N<=5*10^5
T<2
K<=10^9

Source

思路

后缀自动机，由于后缀自动机对应每一个字串，那么只需要记录每个字串的$size$，当$T=0$时为$1$，当$T=1$时为他的$parent$树上的子树大小（注意在这种情况下只能计算插入字符中的$np$，而不能将$nq$计算在内）。

代码

#include <cstdio>
#include <cstring>

const int maxn=500000;
const int maxk=1000000;

struct node
{
  node* tr[26];
  node* par;
  int maxl,size,sum;

  node()
  {
    memset(tr,0,sizeof tr);
    par=NULL;
    maxl=size=sum=0;
  }
};

struct suffix_automaton
{
  node samnode[maxk+10];
  node* pps[maxk+10];
  node* qs;
  node* qlast;
  int cnt,tmp[maxk+10];

  suffix_automaton()
  {
    delete qs;
    delete qlast;
    qlast=qs=new node();
    cnt=0;
  }

  inline int addchr(int ch)
  {
    node* p=qlast;
    node* np=samnode+ ++cnt;
    np->size=1;
    np->maxl=qlast->maxl+1;
    qlast=np;
    while((p!=NULL)&&(p->tr[ch]==NULL))
      {
        p->tr[ch]=np;
        p=p->par;
      }
    if(p==NULL)
      {
        np->par=qs;
        return 0;
      }
    node* q=p->tr[ch];
    if(q->maxl!=p->maxl+1)
      {
        node* nq=samnode+ ++cnt;
        nq->size=0;
        nq->maxl=p->maxl+1;
        memcpy(nq->tr,q->tr,sizeof q->tr);
        nq->par=q->par;
        q->par=np->par=nq;
        while((p!=NULL)&&(p->tr[ch]==q))
          {
            p->tr[ch]=nq;
            p=p->par;
          }
      }
    else
      {
        np->par=q;
      }
    return 0;
  }

  inline int sort_maxl()
  {
    for(register int i=1; i<=cnt; ++i)
      {
        ++tmp[samnode[i].maxl];
      }
    for(register int i=2; i<=cnt; ++i)
      {
        tmp[i]+=tmp[i-1];
      }
    for(register int i=cnt; i; --i)
      {
        pps[tmp[samnode[i].maxl]--]=samnode+i;
      }
    return 0;
  }

  inline int calc_size(int t)
  {
    if(t)
      {
        for(register int i=cnt; i; --i)
          {
            if(pps[i]->par!=qs)
              {
                pps[i]->par->size+=pps[i]->size;
              }
          }
      }
    else
      {
        for(register int i=1; i<=cnt; ++i)
          {
            samnode[i].size=1;
          }
      }
    for(register int i=cnt; i; --i)
      {
        pps[i]->sum=pps[i]->size;
        for(register int j=0; j<26; ++j)
          {
            if(pps[i]->tr[j]!=NULL)
              {
                pps[i]->sum+=pps[i]->tr[j]->sum;
              }
          }
      }
    return 0;
  }

  inline int dfs(node* p,int k)
  {
    if(p->size>=k)
      {
        return 0;
      }
    k-=p->size;
    for(register int i=0; i<26; ++i)
      {
        if(p->tr[i]!=NULL)
          {
            if(p->tr[i]->sum>=k)
              {
                putchar(i+'a');
                dfs(p->tr[i],k);
                return 0;
              }
            k-=p->tr[i]->sum;
          }
      }
    return 0;
  }

  inline int check(int k)
  {
    for(register int i=0; i<26; ++i)
      {
        if(qs->tr[i]!=NULL)
          {
            k-=qs->tr[i]->sum;
          }
      }
    if(k>0)
      {
        return 1;
      }
    return 0;
  }
};

suffix_automaton sam;
int n,k,t;
char s[maxn+10];

int main()
{
  scanf("%s",s+1);
  n=strlen(s+1);
  for(register int i=1; i<=n; ++i)
    {
      sam.addchr(s[i]-'a');
    }
  sam.sort_maxl();
  scanf("%d%d",&t,&k);
  sam.calc_size(t);
  if(sam.check(k))
    {
      puts("-1");
      return 0;
    }
  sam.dfs(sam.qs,k);
  return 0;
}