🦍
ゴリラでもわかるSA-IS

2024/10/14に公開
Rust
 初めにSA-IS[1]は接尾字配列を線形時間[2]で構築できるアルゴリズムの1つです。他の線形時間アルゴリズムと比較して高速かつ省メモリであり、コードもコンパクトです。

 SA-KASA-ISはSA-KA[3]を改良したものです。そこで、SA-ISの解説に入る前にSA-KAを概観します。なお、SA-KAという名称は一般的ではないかもしれません。
以下では、文字列Sのi文字目をS[i]と表記します。また、部分文字列S[i]..S[j]をS[i..j]と略記します。とくに、接尾字はS[i..n]です。
Sは辞書順最小の（ユニークな）番兵\#を末尾に持ちます。これは証明の見通しをよくするためです。実装上は必要ありません。
定義1（L型とS型）　辞書順でS[i..n]\gtrless S[i + 1..n]のとき、S[i..n]およびS[i]をL型・S型と定義します。番兵は常にS型とします。すべての接尾字はL型かS型に分類されます。また、接尾辞の長さが異なるため、等号が成立することはありません。
補題1　線形時間でL型・S型を決定できる。

証明1　Sを右から左に走査する。(1)S[i]=S[i+1]ならばS[i]とs[i+1]は同じ型。(2)S[i]\gtrless S[i+1]ならばS[i]はそれぞれL型・S型。
補題2　L型・S型の接尾辞の一方がソート済みであれば、他方も線形時間でソートできる。

証明2　S型の接尾辞がソートされている場合、後述のアルゴリズム1でソートできる。L型の接尾辞がソートされている場合も同様。
これが、SA-ISおよびSA-KAで肝となる補題です。これにより、L型・S型の接尾辞の一方のみを線形時間で整列できればよいことが分かりました。
補題3　頭文字が同じ接尾辞について、S型がL型よりも辞書順で大きい

証明3　番兵は1度しか登場しないことに注意します。S[i]=S[j]を満足するS型のS[i]とL型のS[j]を考えます。このとき、S[i]\le S[i+1]かつS[j]\ge S[j+1]よりS[i..n]>S[j..n]が成り立ちます。
接尾字配列において接尾辞が頭文字ごとにまとまっていることは自明です。このまとまりをバケットと呼ぶことにします。補題3はバケットの左側にL型が右側にS型が集まっていることを意味します。
アルゴリズム1　接尾字配列をSAとします。(1)ソート済みのS型接尾辞を右から左に走査し、SAのバケットに右から詰めていきます。とくに、SA[1]は番兵です。(2)SAを左から右にたどります。SA[i]が必ず存在します。S[SA[i]-1]がL型であれば、バケットの左から詰めていきます。

証明2（S型からL型）　小さな接尾辞に頭文字を追加してL型の接尾字を構築しています。SAを左から走査することでバケット内で小さなものから順に作られます。L型の接尾字SA[i]から頭文字を除いた部分SA[j]について、i>jが成り立ちます。つまり、SA[j]はソート済みです。形式的な証明は原著論文[3:1]を参照してください。
もう少し丁寧な説明SA[1]=nは番兵\#です。番兵の次に小さな文字を\alphaとします。\alphaから始まる接尾字のうちL型のものは\alpha\dots\alpha\#の形に限られます。したがって、SAを左から走査することでバケット\alphaはソートされます。同様に、\alphaの次に小さな文字\betaから始まる接尾字は\beta\dots\beta\#または\beta\dots\beta\alpha\dots\#のみがL型です。前者が辞書順でバケットに詰められ、後者がそれに続きます。
残る問題はL型・S型の接尾字の一方を線形時間でソートする方法ですが、これはSA-ISで解説することにします。基本的な考え方は同じですが、実装がより簡単だからです。

 SA-ISアルゴリズム1ではすべてのS型接尾字がソート済みであることを要求していますが、実際に使用されているのはL型の右にあるS型のみです。したがって、すべてのS型をソートする必要はありません。
定義2（LMS型）　S型の文字S[i]のうちS[i-1]がL型のものを、とくにLMS型[4]と定義する。
アルゴリズム2（induced sort）　(1)LMS型を大きいものから順にバケットの右から詰める。(2)SAを左から右に走査する。SA[i]が存在し、S[SA[i]-1]がL型であれば、バケットに左から詰める。(3)SAを右から左に走査する。SA[i]が存在し、S[SA[i]+1]が（LMS型を含む）S型であれば、バケットに右から詰める。すでにLMS型がある場合は上書きする。
補題4　LMS型接尾字がソート済みであれば、L型接尾辞を線形時間でソートでき、S型接尾辞も線形時間でソートできる。

証明4　補題2とアルゴリズム2より。
残る問題はLMS型の接尾字を線形時間でソートすることです。LMS型の接尾字は高々\lceil n/2 \rceil個あるので、愚直にソートする訳にはいきません。そこで、LMS型接尾字を比較できるように文字列を圧縮することを考えます。
定義3（LMS部分文字列）　LMS部分文字列を次のように定義する。(1)S[i]とS[j]のみがLMS型であるようなS[i..j]\ (i<j)。(2)番兵そのもの。
定義4（LMS部分文字列）の順序　辞書順に加えて文字の型も考慮する。同じ位置に同じ文字があるとき、それがS型の方がL型よりも小さいとするのが自然。
補題5　圧縮後の文字列をS'とおくと、S[i..n]\gtrless S[j..n]\Leftrightarrow S'[i'..n']\gtrless S'[j'..n']。

証明5　\big|S'[i']\big|=\big|S'[j']\big|のとき、S'[i'+1]とS'[j'+1]を比較すればよい。\big|S'[i']\big|<\big|S'[j']\big|のとき、S'[i']の末尾がS型なので、\big|S'[i']\big|文字目までに大小関係が確定する。
補題6　LMS型をソートせずにinduced sortを行うと、LMS部分文字列がソートされる。

証明6　LMS型接尾字は頭文字についてのみソートされる。したがって、induced sortにより、各接尾辞の頭文字から2文字目以降で最初のLMS型の文字までソートされる。これはLMS部分文字列を含んでいる。
補題5より、LMS部分文字列が相異なる場合、LMS部分文字列をソートすることとLMS型接尾辞をソートすることが同値であることが分かりました。しかしながら、同じLMS部分文字列が複数回登場する場合、定義4からそれらの正しい順序を決めることはできません。
補題7　SA-ISを再帰的に適用することで、S'を正しくソートできる。

証明7　n'<\lceil n/2\rceilより、再帰は停止する。
命題8　SA-ISにより、線形時間で接尾辞配列を構築できる。

証明8　以上の議論による。とくに、補題7より再帰全体の最悪計算量はn+n'+n''+\cdots=n+n/2+n/4+\cdots<2n。

 実装例Rust言語による番兵を用いない実装です。
max_letterは登場する文字のうち最大のものを指します。これはバケットの両端を計算するために必要です。LMS文字列が高々\lceil n/2\rceil種類あるためこのような実装になっています。
実際に使用する場合、strを文字コードなどの「文字の大小関係を維持したユニークなインデックスの列」として与える必要があります。最大のインデックスがmax_letter[5][6]です。
fn sa_is(str: &[usize], max_letter: usize) -> Vec<usize> {
    let n = str.len();

    match n {
        0 => return vec![],
        1 => return vec![0],
        _ => (),
    }

    let (is_s_type, partition_right) = {
        let mut is_s_type = vec![false; n]; // end with L-type
        let mut count_by_initial = vec![0; max_letter + 1];
        for i in (0..n).rev().skip(1) {
            is_s_type[i] = match str[i].cmp(&str[i + 1]) {
                std::cmp::Ordering::Less => true,
                std::cmp::Ordering::Equal => is_s_type[i + 1],
                std::cmp::Ordering::Greater => false,
            };

            count_by_initial[str[i]] += 1;
        }
        count_by_initial[str[n - 1]] += 1;

        for i in 0..max_letter {
            count_by_initial[i + 1] += count_by_initial[i];
        }

        (is_s_type, count_by_initial)
    };

    let induced_sort = |mut sa: Vec<usize>, lms: &[usize]| {
        sa.fill(usize::MAX); // >= n

        // step 1. put LMS-type
        {
            let mut partition_right = partition_right.clone();
            for i in lms.iter().cloned().rev() {
                let initial = str[i];
                partition_right[initial] -= 1;
                sa[partition_right[initial]] = i;
            }
        }

        // step 2. put L-type
        {
            let mut partition_left = vec![0; max_letter + 1];
            for i in 0..max_letter {
                partition_left[i + 1] += partition_right[i];
            }
            // last letter ("phantom" sentinel is always LMS-type and put at the head of 'sa')
            let initial = str[n - 1];
            sa[partition_left[initial]] = n - 1;
            partition_left[initial] += 1;
            // other letters
            for i in 0..n {
                if sa[i] < n && sa[i] > 0 && !is_s_type[sa[i] - 1] {
                    let initial = str[sa[i] - 1];
                    sa[partition_left[initial]] = sa[i] - 1;
                    partition_left[initial] += 1;
                }
            }
        }

        // step 3. put S-type
        {
            let mut partition_right = partition_right.clone();
            for i in (0..n).rev() {
                if sa[i] < n && sa[i] > 0 && is_s_type[sa[i] - 1] {
                    let initial = str[sa[i] - 1];
                    partition_right[initial] -= 1;
                    sa[partition_right[initial]] = sa[i] - 1;
                }
            }
        }

        sa
    };

    // sort LMS-substrings
    let (lms_map, lms) = {
        let mut lms_map = vec![usize::MAX; n];
        let mut lms = Vec::with_capacity(n / 2);
        for i in 1..n {
            if !is_s_type[i - 1] && is_s_type[i] {
                lms_map[i] = lms.len();
                lms.push(i);
            }
        }

        (lms_map, lms)
    };
    let sa = induced_sort(vec![usize::MAX; n], &lms);

    // if all sorted LMS-substrings are unique, LMS-suffixes are also sorted
    if lms.is_empty() {
        return sa;
    }

    let sorted_lms = Vec::from_iter(sa.iter().filter(|&&i| lms_map[i] != usize::MAX).cloned());
    let mut new_max_letter = 0;
    let mut new_str = vec![0; sorted_lms.len()];
    for i in 1..sorted_lms.len() {
        let (l, r) = (sorted_lms[i - 1], sorted_lms[i]);

        // compare adjacent LMS-substrings by length, and if same, then compare by letter and type
        let len_l = lms.get(lms_map[l] + 1).cloned().unwrap_or(n - 1) - l;
        let len_r = lms.get(lms_map[r] + 1).cloned().unwrap_or(n - 1) - r;
        if len_l != len_r
            || (0..=len_l).any(|i| str[l + i] != str[r + i] || is_s_type[l + i] != is_s_type[r + i])
        {
            new_max_letter += 1;
        }

        new_str[lms_map[r]] = new_max_letter;
    }

    // induce suffix array from sorted LMS-suffixes
    let mut sorted_lms = sorted_lms;
    if new_max_letter + 1 < sorted_lms.len() {
        let new_sa = sa_is(&new_str, new_max_letter);
        sorted_lms = Vec::from_iter(new_sa.into_iter().map(|i| lms[i]));
    }

    induced_sort(sa, &sorted_lms)
}

脚注
Ge Nong, Sen Zhang, & Wai Hong Chan. Two Efficient Algorithms for Linear Time Suffix Array Construction. IEEE Trans. Comput. 60, 1471–1484 (2011). ↩︎
座標圧縮されている文字列の場合。実際はO(n \log n)の前処理が必要。考慮する文字の種類を\SigmaとしてO(n + \Sigma)で計算しても良い。 ↩︎
Ko, P. & Aluru, S. Space Efficient Linear Time Construction of Suﬃx Arrays. ↩︎ ↩︎
left-most S-type（もっとも左のS型）の略。S型のかたまりの左端。 ↩︎
max_letter = 0x10FFFFとすればUnicodeで表現できるすべての文字に対応できます。 ↩︎
BWTの高速化などであれば、sa_is(str.as_bytes(), u8::MAX)でよいでしょう。 ↩︎
初めに

SA-KA

SA-IS

実装例

Discussion
