Kesinの知見置き場

知見を共有していきたいじゃないですか

SuffixArrayを作ってみた(Python)

前回の続き
自然言語処理はじめました - Ngramを数え上げまくるを見てSuffixArrayを使ったNgram取得をやってみたくなりました!

前回SuffixArray構築に必須の文字列のソートアルゴリズムのマルチキークイックソートが実装できたのでいよいよSuffixArrayを実装します。
SuffixArrayの構築や、これを使用した文字列の検索は単純な実装なら簡単です。自分はこんな感じに作りました

  • SuffixArray構築
    • 文字列内の各文字のポインタを保存(Pythonではポインタがないので要素番号)
    • ポインタの位置から後ろの文字列でソートしてポインタを入れ替える
  • 部分文字列の検索
    • SuffixArray内から二分探索で検索文字列を検索
    • 発見できたらその前や後ろにも検索文字列が存在しないかチェック
    • 発見できた部分文字列の開始ポインタ位置を全て出力
  • Ngram数え上げ
    • SuffixArrayを端からngram見ていき、同じ文字列だったらカウント+。違う文字列になったら出力
#coding: utf-8
'''
単純なSuffixArrayの実装
部分文字列の検索
Ngram数え上げ
'''
def multikeysort(string, indices, cmpindex=0):
    """単純なマルチキークイックソート
    文字列のインデックスを使用したソート"""
    if not indices:
        return []
    pivot = indices[0]
    left=[]
    mid=[]
    right=[]
    #cmpindex番目の文字を比較
    for index in (indices[i] for i in xrange(1, len(indices))):
        if string[index+cmpindex] < string[pivot+cmpindex]:
            left.append(index)
        elif string[index+cmpindex] == string[pivot+cmpindex]:
            mid.append(index)
        else:
            right.append(index)
    mid.append(pivot)
    #番兵"$"によるIndexError防止と計算量軽減
    if not string[pivot+cmpindex] == '$' and not len(mid) < 2:
        mid = multikeysort(string, mid, cmpindex+1)
    left = multikeysort(string, left, cmpindex)
    right = multikeysort(string, right, cmpindex)
    return left + mid + right

def get_suffixarray(string):
    """文字列からSuffixArrayを構築"""
    indices = xrange(len(string))
    return multikeysort(string, indices)

def binary_search(string, suffixarray, query):
    """部分文字列を見つけたときのmidを返す"""
    endindex = len(query)
    stringlen = len(string)
    min = 0
    max = len(suffixarray)-1
    while True:
        mid = (min+max) / 2
        suffix = suffixarray[mid]
        #IndexError防止
        if not suffix + endindex > stringlen:
            cmpstring = string[suffix:suffix+endindex]
            
        if max < min:
            return -1
        elif cmpstring < query:
            min = mid + 1
        elif cmpstring > query:
            max = mid - 1
        #elif string[suffix][:endindex] == query:
        else:
            return mid
def search_all(string, suffixarray, query):
    """部分文字列の位置を全て検索
    見つかれば接尾辞要素番号"""
    endindex = len(query)
    mid = binary_search(string, suffixarray, query)
    left = mid
    right = mid + 1
    while True:
        suffix = suffixarray[left]
        cmpstring = string[suffix:suffix+endindex]
        if cmpstring == query:
            yield suffix
            left -= 1
        else:
            break
    while True:
        suffix = suffixarray[right]
        cmpstring = string[suffix:suffix+endindex]
        if cmpstring == query:
            yield suffix
            right += 1
        else:
            break
def count_ngram(string, suffixarray, n):
    """n-gramを数えるイテレータ"""
    #SuffixArrayからngram以上の文字列を取得するジェネレータ。"$"を除くので-1
    cmp_strings = (string[i:i+n] for i in suffixarray if not i+n > len(string)-1)
    #初期化
    count_string = cmp_strings.next()
    count = 1
    for cmp_string in cmp_strings:
        if count_string == cmp_string:
            count += 1
        else:
            yield (count_string, count)
            count_string = cmp_string
            count = 1
    yield (count_string, count)

#SuffixArrayの構築
>>> string = 'abracadabra$'
>>> suffixarray = get_suffixarray(string)
>>> for suffix in [string[i:] for i in suffixarray]:
...     print suffix
$
a$
abra$
abracadabra$
acadabra$
adabra$
bra$
bracadabra$
cadabra$
dabra$
ra$
racadabra$

>>> string = u"イカちゃんかわいいイカちゃん" + "$"
>>> suffixarray = get_suffixarray(string)
>>> for suffix in [string[i:] for i in suffixarray]:
...     print suffix
... 
$
いいイカちゃん$
いイカちゃん$
かわいいイカちゃん$
ちゃん$
ちゃんかわいいイカちゃん$
ゃん$
ゃんかわいいイカちゃん$
わいいイカちゃん$
ん$
んかわいいイカちゃん$
イカちゃん$
イカちゃんかわいいイカちゃん$
カちゃん$
カちゃんかわいいイカちゃん$

#部分文字列の検索
>>> query = u'イカ'
>>> print [i for i in search_all(string, suffixarray, query)]
[9, 0]

#Ngram数え上げ
>>> for word, n in count_ngram(string, suffixarray, 2):
...     print word, n
... 
いい 1
いイ 1
かわ 1
ちゃ 2
ゃん 2
わい 1
んか 1
イカ 2
カち 2

できたことにはできたけど、検索とngram数え上げの部分がかなり見苦しい・・・。番兵を上手く使えよって話なんですが、一文字ずつ比較すると遅くなりそうな気がしてしまってこんな感じにせざるを得なかったです。他にいい方法ないかな。

より実用的なものを目指すのであればSuffixArrayの圧縮や、効率的なソート、検索手法を組み合わせる必要があるかと思います。ちょっと調べたんですが、複雑すぎてサッパリ(-_-;)突き詰めるならデータ構造について知識がないとダメですね。