Suffix Array の解説と実装方法[C#][C++]

Suffix Array(接尾辞配列) とは

接尾辞配列 – Wikipedia

接尾辞配列（せつびじはいれつ）やサフィックス・アレイ（英: suffix array）とは、文字列の接尾辞（開始位置を異にし終端位置を元の文字列と同じくする部分文字列）の文字列中の開始位置を要素とする配列を、接尾辞に関して辞書順に並べ替えて得られる配列である。

Suffix Array を構築することで文字列検索を高速に行うことができます。

例えばある文字列Sに文字列Tが含まれるかどうかを判定する場合、何も考えずに全探索すると計算量 O(|S||T|) かかります。Suffix Array がすでに構築されている場合、二分探索で判定できるので計算量が O(log|S|) で済むようになります。

Suffix Array の具体例

では Suffix Array の具体例を見てみます。文字通りSuffix(接尾辞)の配列です。Suffixは文字列から先頭x文字を取り除いた文字列になります。

例えば "abracadabra" という11文字の Suffix は先頭から0~11文字取り除けるので以下の12種類が Suffix となります。

先頭から全文字取り除いた空文字列もSuffixとして扱います。開始位置は元の文字列のどの位置から開始したSuffixかを意味しています。

この状態ですべての開始位置から始まる末尾までの文字列(Suffix)が列挙されています。

このままの状態では何の意味もありません。このSuffixをソートして初めて意味のあるSuffix Arrayになります。ということでSuffixでソートして得られるのが以下の配列です。

Suffix Arrayは実際に文字列として配列は持つ必要はなく、開始位置の配列を保持することになります。このようなSuffix Arrayがあれば部分文字列の検索が高速に行えます。

例えば、"bra" という文字列を含むかどうか判定するには上記のようなSuffix Arrayを構築し、二分探索してやれば存在するか判定できますし、その開始位置もわかります。さらに複数含む場合そのすべての開始位置も二分探索で列挙できます。つまり、文字列Sが文字列Tを含むかどうかの判定は、計算量 O(MlogN) で行えるということです。(M=|T|, N=|S|)

ただし問題点は Suffix のソートにあります。通常のソート(クイックソートなど)を使って愚直にソートすると、N=|S|としたとき、計算量 O(N2logN) かかります。ソートの計算量が O(NlogN) で文字列の比較1回につき O(N) かかるためです。したがって通常のソートよりも高速なアルゴリズムを使ってソートを行わなければなりません。

Suffix Array構築部分で O(N2logN) かかってしまっては、全探索による部分文字列の判定よりも遅くなってしまいます。

Suffix Array による部分文字列の検索を実装してみる

ひとまずソート部分は後で考えるとして、部分文字列を愚直にソートする方法で実装してみます。C# だと以下のように実装できます。

class SuffixArray
{
    private string S;
    private int N;
    public int[] SA;
    public SuffixArray(string s)
    {
        this.S = s;
        this.N = s.Length;
        this.SA = new int[N + 1];
        CreateSuffixArray();
    }
    private void CreateSuffixArray()
    {
        for (int i = 0; i < N + 1; i++)
        {
            SA[i] = i;
        }
        // ソート ※計算量O(N^2logN)で遅い
        Array.Sort(SA, CompareSA);
    }
    public int CompareSA(int i, int j)
    {
        // 単純に開始位置から末尾までの文字列で比較
        return string.CompareOrdinal(S, i, S, j, N);
    }

    // 部分文字列にTが含まれるかどうかを判定する
    public bool Contains(string t)
    {
        var l = 0;
        var r = N;
        while (r - l > 1)
        {
            var mid = l + (r - l) / 2;
            var index = SA[mid];
            var cmp = string.CompareOrdinal(S, index, t, 0, t.Length);
            if (cmp < 0) l = mid;
            else r = mid;
        }
        return string.CompareOrdinal(S, SA[r], t, 0, t.Length) == 0;
    }
}

とりあえず部分文字列の検索処理を実装しました。含まれるかどうかを判定するだけなので答えはbool値にしています。

var s = "abracadabra";
var sa = new SuffixArray(s);
Console.WriteLine(sa.Contains("ab"));          // True
Console.WriteLine(sa.Contains("dabra"));       // True
Console.WriteLine(sa.Contains("adr"));         // False
Console.WriteLine(sa.Contains("ada"));         // True
Console.WriteLine(sa.Contains("abracadabra")); // True

こんな感じで部分文字列の検索が実行できることがわかります。では計算量のボトルネックになっているソート部分を改良してみます。

ダブリングのアイデアを使ったソートの高速化

ソートを高速化するにあたって使用するアイデアはダブリングです。

ダブリングはあらかじめ1つ次の要素についての計算を行っておき、1つ次の要素のさらに1つ次の要素についての計算を即座に得ようという考えです。例えば繰り返し二乗法がダブリングの1種です。

[C#] 繰り返し二乗法で n^p % mod を拘束に求める │ Web備忘録

N^p を求めるのに N*N*N*...*N みたいに繰り返すのではなく、

• N*N=N²
• N² * N² = N⁴
• N⁴ * N⁴ = N⁸
• …
• N^k * N^k = N^2k

みたいに倍々ゲームで計算を行うことができれば、計算量のオーダーがlogになります。これがダブリングの考え方です。ではこのアイデアをどのようにソートに使うかを考えます。

Suffix の先頭k文字の大小関係

ダブリングを使うことで各Suffix先頭k文字の大小がわかると、先頭2k文字の大小関係が判定できます。これを説明します。

Suffix各位置から1文字の大小関係を求めておきます。i番目のSuffixの先頭k文字の値の大きさのことを Rankk[i] と呼ぶことにします。

Rank1[i] はすぐに求まります。各Suffixの先頭1文字についての文字コードがあればよいです。

これは各Suffix先頭1文字の大小関係を比較するための数値です。単純に先頭1文字のAsciiコードです。空文字については-1としてどの文字よりも前に来るようにします。わかりやすくするために文字コードにしてますが、大小関係の判定ができればよいだけなので実装ではもっと小さい数字を使うようになります。

さて、この大小関係がわかった状態から先頭2文字についての大小関係を考えます。

まず、先頭1文字についての大小関係は Rank1[i] が すでに求まっているのでこれを使って判定すればよいです。Rank1[i] が異なるときは先頭1文字が異なるのでこの大小関係が2文字の大小関係と等しくなります。

Rank1[i]が等しいとき、先頭1文字が等しいので次の1文字についての大小を比較しなければなりません。

ここでSuffixであるということを思い出すと、次の1文字についての1文字についての大きさは Rank1[i + 1] ですでに計算済であることがわかります。つまり開始位置を1ずらした Suffix が存在すれば、そのSuffixの先頭1文字の大きさが、次の1文字の大きさになります。

こんな感じで先頭2文字についてのソートを行う場合は各文字列の比較を Rank1[i] と Rank1[i+1] のペアについて行います。こうすることで比較自体がO(1)で済むようになります。先頭2文字の比較が行えるようになればソートもできます。ソートができれば先頭2文字のRankがつけられます。

これを繰り返せば、4文字, 8文字 … 2k文字のソートが行えるようになり、最終的にN文字のソートが完了します。

つまり、Suffix先頭k文字の大小関係を判定するための値 Rankk が求まっていると、Suffix先頭2k文字の比較が行えるためソートが高速化します。つまり、先頭2k文字のソートは Rankk[i] と Rank1[i+k] のペアについて比較を行いソートするということです。

そして元の文字列以上になるまでソートを繰り返すと、logN回のソートを行うことになり、1回のソートにつき O(NlogN) の計算なので、結果 O(Nlog2N) の計算量で Suffix Array のソートが完了します。

この Rank を使ったソート手順については実装を見て、ステップを追いかけながらでないと理解が難しいです。

実装

解説はC#のコードで行いますが、C++と基本変わらないです。

C#

class SuffixArray
{
    private int N;
    private string S;
    private int[] Rank;
    private int[] Tmp;
    public int[] SA;
    private int K;
    public SuffixArray(string s)
    {
        S = s;
        N = S.Length;
        Rank = new int[N + 1];
        SA = new int[N + 1];
        Tmp = new int[N + 1];
        CreateSuffixArray();
    }
    private void CreateSuffixArray()
    {
        // ランクの初期値Rank1は文字コードとする
        for (int i = 0; i <= N; i++)
        {
            SA[i] = i;
            Rank[i] = i < N ? S[i] : -1;
        }

        // K文字のRankが確定してソート済の状態から
        // 2k文字のソートを行う
        for (K = 1; K <= N; K *= 2)
        {
            // 2k文字のソートを行う
            Array.Sort(SA, CompareSA);

            // 2k文字のランクのペアで次のランク(2k)を作る
            Tmp[SA[0]] = 0;
            for (int i = 1; i <= N; i++)
            {
                Tmp[SA[i]] = Tmp[SA[i - 1]]
                    + (CompareSA(SA[i - 1], SA[i]) < 0 ? 1 : 0);
            }
            // 作成した次のランク(2k)に更新する
            for (int i = 0; i <= N; i++)
            {
                Rank[i] = Tmp[i];
            }
        }
    }
    // 接尾辞配列のソート用比較関数(開始位置がi, jのSuffixを比較)
    // (rank[i], rank[i+k])と(rank[j], rank[j+k])を比較する
    private int CompareSA(int i, int j)
    {
        // 先頭K文字が異なるため比較が終了する
        if (Rank[i] != Rank[j]) return Rank[i] - Rank[j];
        else
        {
            // 次のK文字のランクで比較する
            // 次のK文字が存在しない場合は -1 で小さくなるようにする
            var ranki = i + K <= N ? Rank[i + K] : -1;
            var rankj = j + K <= N ? Rank[j + K] : -1;
            return ranki - rankj;
        }
    }

    // 部分文字列にTが含まれるかどうかを判定する
    public bool Contains(string t)
    {
        var l = 0;
        var r = N;
        while (r - l > 1)
        {
            var mid = l + (r - l) / 2;
            var index = SA[mid];
            var cmp = string.CompareOrdinal(S, index, t, 0, t.Length);
            if (cmp < 0) l = mid;
            else r = mid;
        }
        return string.CompareOrdinal(S, SA[r], t, 0, t.Length) == 0;
    }
}

実装は蟻本準拠です。

C++

struct SuffixArray {
  string s;
  int n, k;
  vector<int> sa, tmp, rank_sa;

  struct CompareSA {
    int n, k;
    const vector<int> &rank;
    CompareSA(int n, int k, const vector<int> &rank_sa)
        : n(n), k(k), rank(rank_sa) {
    }
    bool operator()(int i, int j) {
      if(rank[i] != rank[j])
        return rank[i] < rank[j];
      else {
        int rank_ik = (i + k <= n ? rank[i + k] : -1);
        int rank_jk = (j + k <= n ? rank[j + k] : -1);
        return rank_ik < rank_jk;
      }
    }
  };
  SuffixArray(const string &S) {
    s = S;
    n = S.size();
    sa = vector<int>(n + 1);
    rank_sa = vector<int>(n + 1);
    tmp = vector<int>(n + 1);
    createSuffixArray();
  }
  void createSuffixArray() {
    for(int i = 0; i <= n; i++) {
      sa[i] = i;
      rank_sa[i] = i < n ? (int)s[i] : -1;
    }

    for(k = 1; k <= n; k *= 2) {
      CompareSA csa(n, k, rank_sa);
      sort(sa.begin(), sa.end(), csa);

      tmp[sa[0]] = 0;
      for(int i = 1; i <= n; i++) {
        tmp[sa[i]] = tmp[sa[i - 1]];
        if(csa(sa[i - 1], sa[i])) tmp[sa[i]]++;
      }
      for(int i = 0; i <= n; i++) {
        rank_sa[i] = tmp[i];
      }
    }
  }

  bool contains(string t) {
    int l = 0;
    int r = n;
    while(r - l > 1) {
      int mid = l + (r - l) / 2;
      int index = sa[mid];
      int cmp = s.compare(index, t.size(), t);
      if(cmp < 0)
        l = mid;
      else
        r = mid;
    }
    return s.compare(sa[r], t.size(), t) == 0;
  }
};

比較関数がそのままだと使えなかったのでちょっと変えましたが大枠は変わりません。

コード解説

C# の実装を解説します。

CreateSuffixArray() で Suffix Array を作成します。

まずあらかじめ先頭1文字のランクを設定しておきます。これは先頭1文字の文字コードを採用します。

// ランクの初期値Rank1は文字コードとする
for (int i = 0; i <= N; i++)
{
    SA[i] = i;
    Rank[i] = i < N ? S[i] : -1;
}

とりあえず上のようなデータが初期値になります。この状態からソートを行っていきます。

// K文字のRankが確定してソート済の状態から
// 2k文字のソートを行う
for (K = 1; K <= N; K *= 2)
{
    // 2k文字のソートを行う
    Array.Sort(SA, CompareSA);

    // 2k文字のランクのペアで次のランク(2k)を作る
    Tmp[SA[0]] = 0;
    for (int i = 1; i <= N; i++)
    {
        Tmp[SA[i]] = Tmp[SA[i - 1]]
            + (CompareSA(SA[i - 1], SA[i]) < 0 ? 1 : 0);
    }
    // 作成した次のランク(2k)に更新する
    for (int i = 0; i <= N; i++)
    {
        Rank[i] = Tmp[i];
    }
}

2のべき乗でソートを繰り返します。2文字, 4文字, 8文字…2k文字(2k>=N)のソート順番に実行します。先頭2K文字でソートを実行した後は2K文字のランクに更新し、次のループに移ります。

// 2k文字のソートを行う
Array.Sort(SA, CompareSA);

この部分で Suffix Array のソートを実行しています。この時、Rankkがすでに求まっています。

比較関数

ソート時の比較関数(CompareSA)は以下のようになります。

// 接尾辞配列のソート用比較関数(開始位置がi, jのSuffixを比較)
// (rank[i], rank[i+k])と(rank[j], rank[j+k])を比較する
private int CompareSA(int i, int j)
{
    // 先頭K文字が異なるため比較が終了する
    if (Rank[i] != Rank[j]) return Rank[i] - Rank[j];
    else
    {
        // 次のK文字のランクで比較する
        // 次のK文字が存在しない場合は -1 で小さくなるようにする
        var ranki = i + K <= N ? Rank[i + K] : -1;
        var rankj = j + K <= N ? Rank[j + K] : -1;
        return ranki - rankj;
    }
}

先頭K文字のRankが異なる場合は次のK文字を見る必要がないので、先頭K文字の比較結果を返すようにします。

先頭K文字が等しいときは次のK文字について比較して確認します。次のK文字は Rank2k[i+k] で求まります。なお、i+K が N より大きいとき、元の文字列Sの末尾をはみ出しているので、-1とします。つまり、次のK文字は存在しません。

ランクの更新

先頭2K文字で Suffix Array がソート済の時、K文字のランクから2K文字のランクを生成するには次のようにします。

// 2k文字のランクのペアで次のランク(2k)を作る
Tmp[SA[0]] = 0;
for (int i = 1; i <= N; i++)
{
    Tmp[SA[i]] = Tmp[SA[i - 1]]
        + (CompareSA(SA[i - 1], SA[i]) < 0 ? 1 : 0);
}
// 作成した次のランク(2k)に更新する
for (int i = 0; i <= N; i++)
{
    Rank[i] = Tmp[i];
}

SA[0] は空文字列が入っていて、これが常に一番先頭に来るのでこれをランク0とします。この0を基準に、2K文字のランクを生成していきます。

例えば先頭2文字でソートした後の、Rank₂ を生成すると次のようになります。

ランクはあくまでこの表における相対的な比較順位なのでこのようになります。先頭2文字が等しいデータについてはランクが同等になっています。例えば "abracadabra", "abra" は先頭2文字が "ab" で同じなので、ランクが同じになります。

この表のデータが得られると、4文字でのソートが可能になります。

先頭4文字でソートする際の比較の例として、"bracadabra" と "bra" は上の表だと順位が正しくありませんが、4文字で比較する場合は、 (Rank2[1], Rank2[1+2]) と (Rank2[8], Rank2[8+2]) の比較を行います。

これは上のデータを確認すると (5, 3) と (5, 1) の比較になり、"bra" のほうが小さいということが判定できます。

"abracadabra" の Suffix Array 構築時の遷移

SA と Rank の状態の遷移を出力してみます。以下のような遷移でソートが実行されていきます。最終的に2K文字でのソート後、2K>=N になっていればソートが完了しています。

初期状態(先頭1文字のランクが設定されている)

先頭2文字でソート後

先頭4文字でソート後

先頭8文字でソート後

先頭16文字でソート後

テスト

では正しく実装できているか、テスト代わりにAOJの問題を解いてみます。

judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ALDS1_14_D&lang=ja

文字列 T の中に文字列 P が出現するかどうかを判定してください。最初に T が与えられたうえで、質問として Q 個の Pi が与えられます。

制約は以下のようになっています。

• 1 ≤ T の長さ ≤ 1000000
• 1 ≤ Pi の長さ ≤ 1000
• 1 ≤ Q ≤ 10000
• 文字列は、英小文字、英大文字、数字のみで構成されている。

ナイーブな実装だと O(|T||Pi|Q) かかるところ、Suffix Array を使うことで、O(|T|log2|T|) で Suffix Array を構築でき、クエリには O(|Pi|log|T|) で答えることができるようになります。

ACできたので Suffix Array で部分文字列の検索がうまく実装できていることが確認できました。

このままの検索処理の実装では解けませんが、少し実装を工夫してすべての一致した検索結果を返すようにすると以下の問題も解けるようになります。

• No.430 文字列検索 – yukicoder

以上。