Codeforces 128 B. String

久々にCodeforces参加したらボロボロだった。

問題

文字列が与えられる。
与えられる文字列の空でない部分文字列を辞書順に並べ替えたとき、
k番目に来るものを求めよ。

制約条件

文字列の長さ≦100000
k≦100000

方針

自分が立てたのは次のような感じ。

suffix arrayを作った後、1文字ずつ決定していく。
k番目の文字列がxより後のアルファベットから始まるとすると、
suffix arrayのxで始まるの文字列の、全文字数よりもkが大きくなる。

xではじまるならば、kからxで始まるsuffixの数を引く。
suffixのxで始まる部分で、xより後ろの部分を考えて上の操作を繰り返す。

ソースコード

struct SAComp {
  const int h, *g;
  SAComp(int h, int* g) : h(h), g(g) {}
  bool operator() (int a, int b) {
    return a == b ? false : g[a] != g[b] ? g[a] < g[b] : g[a+h] < g[b+h];
  }
};
int *buildSA(const char* t, int n) {
  int g[n+1], b[n+1], *v = new int[n+1];
  rep(i,n+1) v[i] = i, g[i] = t[i];
  b[0] = 0; b[n] = 0;

  sort(v, v+n+1, SAComp(0, g));
  for(int h = 1; b[n] != n ; h *= 2) {
    SAComp comp(h, g);
    sort(v, v+n+1, comp);
    rep(i, n) b[i+1] = b[i] + comp(v[i], v[i+1]);
    rep(i, n+1) g[v[i]] = b[i];
  }
  return v;
}

string in;

void run(){
  cin>>in;
  int n=in.size(),k; cin>>k;
  
  if(k>n*(n+1ll)/2){
    cout<<"No such line."<<endl;
    return;
  }
  
  int *sa=buildSA(in.c_str(),n);
  
  ll *sum=new ll[n+2];
  sum[0]=0;
  rep(i,n+1){
    sum[i+1]=sum[i]+n-sa[i];
  }
  
  string ans;
  
  ll len=0,p=0,q=n,dbg=0;
  while(k>0){
    
    char c;
    int lb,ub;
    for(c='a';c<='z';c++){
      //範囲からlen文字目がcであるような文字の最初と最後を探す
      int lo=p-1,hi=q,mid;
      while(lo+1<hi){
        mid=(lo+hi)/2;
        if(in[sa[mid]+len]<c)lo=mid;
        else hi=mid;
      }
      lb=hi;
      
      lo=p; hi=q+1;
      while(lo+1<hi){
        mid=(lo+hi)/2;
        if(in[sa[mid]+len]<=c)lo=mid;
        else hi=mid;
      }
      ub=lo;
      
      if(lb>ub||sa[lb]+len>=n||in[sa[lb]+len]!=c)continue;

      if(k>sum[ub+1]-sum[lb]-len*(ub-lb+1))k-=sum[ub+1]-sum[lb]-len*(ub-lb+1);
      else break;
    }
    ans.pb(c);
    k-=ub-lb+1;
    p=lb; q=ub;
    len++;
  }
  cout<<ans<<endl;
}