Pythonで全部分字列の和集合から指定した位置の部分文字列を求めるプログラム

問題の概要

n個の文字列 str1、str2、str3、…、strn が与えられているとします。ここで、substri は stri が持つすべての部分文字列を格納する集合であり、すべての substr 集合の和集合を substr_union と定義します。

このとき、q 個のクエリが与えられ、各クエリに対して substr_union の q 番目の要素を求めます。substr_union は辞書順にソートされており、インデックスは 1 から始まるものとします。

具体例

たとえば、文字列のリストが ['pqr', 'pqt']、クエリが [4, 7, 9] の場合、出力は ['pqt', 'qt', 't'] になります。

まず、各文字列から得られる部分文字列を確認しましょう。

• subs_str_1 = {p, pq, pqr, q, qr, r}
• sub_str_2 = {p, pq, pqt, q, qt, t}

これら2つの集合の和集合 substr_union は {p, pq, pqr, pqt, q, qr, qt, r, t} となります。

したがって、インデックス 4、7、9 に対応する要素はそれぞれ 'pqt'、'qt'、't' です。

解決のための手順

この問題を解くには、以下の手順に従います。

• 関数 lng_i() を定義する（引数：suff、lng、i）。ソート済みの接尾辞リストから i 番目の部分文字列を取り出す関数です。
  • d := (suff, lng) をペアにしたタプルのリスト
  • lo := 0、hi := 0 で初期化
  • d 内の各タプル (suf, lng) について次を繰り返す。
    • lng が None の場合は 0 に置き換える
    • hi := hi + (suf の長さ − lng)
    • hi − 1 が i と等しければ suf を返す
    • hi − 1 が i より大きければ、p を lng から suf の長さ未満まで走査し、lo + p が i と一致する位置 j を見つけて suf[0 : j+1] を返す
  • lo := hi と更新して次の接尾辞へ進む
  • 最後まで見つからなければ False を返す
• 関数 hlp_ii() を定義する（引数：str1、str2）。2つの文字列の最長共通接頭辞（LCP）の長さを求める関数です。
  • ub := min(len(str1), len(str2))
  • cnt := 0
  • i を 0 から ub 未満まで走査し、str1[i] と str2[i] が一致すれば cnt を +1、不一致ならその時点で cnt を返す
  • ループが完了すれば cnt を返す
• t_dict := 新しいマップ（辞書）を作成する
• suff := 新しいリストを作成する
• lng := 新しいリストを作成する
• strings 内の各文字列について、すべての接尾辞を生成し、t_dict に未登録のものだけを suff に追加する
• suff をソートする
• suff_len := len(suff) とし、隣接する接尾辞同士のLCP長を hlp_ii() で計算して lng に格納する（先頭要素は None）
• res := 新しいリストを作成する
• q_list 内の各クエリ q について、lng_i(suff, lng, q−1) の結果を res に追加する
• res を返す

実装例

理解を深めるために、以下の実装を見てみましょう。

def lng_i(suff, lng, i):
    d = zip(suff,lng)
    lo = hi = 0
    for suf, lng in d:
        if lng is None:
            lng = 0
        hi += len(suf) - lng
        if hi - 1 == i:
            return suf
        elif hi - 1 > i:
            for p, q in enumerate(list(range(lng, len(suf)))):
                if lo + p == i:
                    return suf[:q+1]
        lo = hi
    return False

def hlp_ii(str1,str2):
    ub = min(len(str1), len(str2))
    cnt = 0
    for i in range(ub):
        if str1[i] == str2[i]:
            cnt += 1
        else:
            return cnt
    return cnt

def solve(strings,q_list):
    t_dict = {}
    suff = []
    lng = []
    for str in strings:
        for i in range(len(str)):
            value = str[i:]
            if value not in t_dict:
                t_dict[value] = 1
                suff.append(value)
    suff.sort()
    suff_len = len(suff)
    for i in range(suff_len):
        if i == 0:
            lng.append(None)
        else:
            lng.append(hlp_ii(suff[i-1], suff[i]))
    res = []
    for q in q_list:
        (res.append(lng_i(suff, lng, q-1)))
    return res

print(solve(['pqr', 'pqt'], [4, 7, 9]))

入力

['pqr', 'pqt'], [4, 7, 9]

出力

['pqt', 'qt', 't']

アルゴリズムのポイント

この手法は接尾辞配列（suffix array）の考え方を応用しています。各文字列のすべての接尾辞を集めてソートすると、隣接する接尾辞どうしの共通接頭辞（LCP）のぶんだけ部分文字列が重複します。そこで「接尾辞の長さ − 直前の接尾辞とのLCP長」を累積していけば、重複を除いた部分文字列の辞書順インデックスを効率よく特定できます。全部分字列を実際に生成してソートする方法に比べ、メモリ使用量と計算量を大幅に抑えられるのが大きな利点です。