Z algorithmで文字列探索を試す

はじめに

名前がかっこいい。
codeforcesにある解説を試してみる。

Z algorithmとは

文字列Sと部分文字列S[i..]の最長共通接頭辞数をZ[i]とし、すべてのiについて、それをO(n)で求めるアルゴリズム
- 単純な方法だとO(n^2)
1996,97年あたりにGusfieldによって提案された

アルゴリズム

うまく過去の情報をメモしておいて計算量を減らす。

iのZ-boxを範囲[i, i+Z[i]-1]と定義する
- あるiについて、i以下のjについてZ-boxを考え、
- Z-boxの終点で最も右側にあるものをr_i、そのr_iのペアとなるZ-boxの始点をl_iとする
Z[k]を求めるときは、過去のZとr_{k-1}とl_{k-1}があれば求められる
k>rの場合
- rを超えてるので、新しくkでのZ-boxを考える必要があるので、部分文字列S[k..]の接頭辞を比較する
- k=l, r=k+Z[k]-1とする
k<=rの場合
- kは計算済みのZ-box内にいるので、その情報l,rが使えるかもしれない
- そこで、Z-boxの[k, r]に対して、文字列Sの接頭部分の[k-l, r-l]の文字列を考える
- Z[k-l]が|r-k+1|より小さいならば、kでのZ-boxはk-lのと同じになるで、Z[k]=Z[k-l]になる
- Z[k-l]が|r-k+1|より大きいならば、kでのZ-boxはk-lのと同じではないので、新しくkでのZ-boxを考える(k>rの時と同様)

Z algorithmを使った文字列探索

文字列Tからパターン文字列Pを探索するために、この2つを連結したP$Tという文字列を作る。
この文字列について、Tの部分のZ[i]を全て求め、Z[i]=|P|となるものを列挙すれば探索できる。

コード

codeforcesのを参考に。

#include <iostream>
#include <algorithm>
#include <string>
#include <vector>

//naive : O(n^2)
int findLCP(const std::string& s){
  int ret = 0;
  int n = s.length();
  std::vector<int> z(n, 0);

  for(int i=1; i<n; i++){
    for(int j=0; i+j<n; j++){
      if(s[j] == s[i+j]) z[i]++;
      else break;
    }
    ret = std::max(ret, z[i]);
  }
  return ret;
}

//Z algorithm : O(n)
int Zalgorithm(const std::string& s){
  int ret = 0;
  int n = s.length();
  std::vector<int> z(n, 0);

  int L = 0, R = 0;
  for(int i=1; i<n; i++){
    if(i>R){
      L = R = i;
      while(R<n && s[R-L] == s[R]) R++;
      z[i] = R-L;
      R--;
    }else{
      int k = i-L;
      if(z[k] < R-i+1){
        z[i] = z[k];
      }else{
        L = i;
        while(R<n && s[R-L] == s[R]) R++;
        z[i] = R-L;
        R--;
      }
    }
    ret = std::max(ret, z[i]);
  }
  return ret;
}

//exact search by Z algorithm
std::vector<int> Zsearch(const std::string& T, const std::string& P){
  std::string S = P + "$" + T;
  std::vector<int> z(S.length(), 0);

  int L = 0, R = 0;
  for(int i=1; i<S.length(); i++){
    if(i>R){
      L = R = i;
      while(R<S.length() && S[R-L] == S[R]) R++;
      z[i] = R-L;
      R--;
    }else{
      int k = i-L;
      if(z[k] < R-i+1){
        z[i] = z[k];
      }else{
        L = i;
        while(R<S.length() && S[R-L] == S[R]) R++;
        z[i] = R-L;
        R--;
      }
    }
  }
  
  std::vector<int> ret;
  for(int i=P.length()+1; i<S.length(); i++){
    if(z[i] == P.length()){
      ret.push_back( i-(P.length()+1) );
    }
  }
  return ret;
}


int main(){
  
  std::string str = "aabadaabcaaba";

  //最長共通接頭辞の長さを返す
  std::cout << findLCP(str) << std::endl;
  std::cout << Zalgorithm(str) << std::endl;

  std::cout << "----" << std::endl;

  //文字列strからパターン文字列patをすべて探す
  std::string pat = "aab";
  std::vector<int> ret = Zsearch(str, pat);
  for(int i=0; i<ret.size(); i++){
    std::cout << ret[i] << std::endl;
  }

  return 0;
}