Pythonで文字列の指定範囲ごとに異なる部分文字列の数を求めるプログラム

長さ n の文字列 s と、クエリのリスト Q が与えられます。各クエリ Q[i] はペア (l, r) を含んでおり、文字列 s のインデックス l から r まで(両端を含む)の範囲に存在する異なる部分文字列の数を数える必要があります。

たとえば、入力が s = "pptpp"、Q = [(1,1), (1,4), (1,1), (0,2)] の場合、出力は [1, 8, 1, 5] になります。その理由は次の通りです。

• クエリ (1, 1) の場合:範囲内の部分文字列は「p」のみなので、出力は 1

• クエリ (1, 4) の場合:部分文字列は「p」「t」「pp」「pt」「tp」「ptp」「tpp」「ptpp」の 8 種類なので、出力は 8

• 再びクエリ (1, 1) の場合:部分文字列は「p」のみなので、出力は 1

• クエリ (0, 2) の場合:部分文字列は「p」「t」「pp」「pt」「ppt」の 5 種類なので、出力は 5

解法のアプローチ

この問題を解くには、接尾辞配列(Suffix Array)と、それに付随する LCP配列(最長共通接頭辞配列) を利用します。LCP配列の構築には Kasaiのアルゴリズム を使います。ポイントとなるのは次の性質です。

異なる部分文字列の総数 = Σ(各接尾辞の長さ − 隣接する接尾辞とのLCPの長さ)

接尾辞を辞書順にソートすると、隣接する接尾辞どうしが共有している接頭辞こそが「二重にカウントされる部分文字列」に対応します。したがって、各接尾辞の長さから LCP を差し引いて合計すれば、重複を除いた部分文字列の数が求まります。

Kasaiのアルゴリズム(LCP配列の構築)

1. サイズ n の配列 lcp と inv を 0 で初期化する
2. i を 0 から n−1 までループし、inv[suff[i]] = i を設定する
3. k = 0 とし、i を 0 から n−1 までループする
   • inv[i] が n−1 なら、k を 0 に戻して次の反復へ進む
   • j = suff[inv[i] + 1] とする
   • s[i+k] と s[j+k] が一致する限り k を増やし続ける
   • lcp[inv[i]] = k を記録し、k > 0 なら k を 1 減らす
4. lcp を返す

クエリへの回答(メイン処理)

1. 結果を格納する空のリスト res を用意する
2. 各クエリ (left, right) について次を行う
   • sub = s[left:right+1] として対象範囲の部分文字列を取り出す(length = right − left + 1)
   • (インデックス, 接尾辞文字列) のペアをすべて作成し、接尾辞文字列の辞書順でソートする
   • ソート済みの接尾辞配列 suff に対して kasai() を呼び出し、LCP 配列を取得する
   • count を最初の接尾辞の長さで初期化し、残りの各接尾辞について「長さ − LCP」を加算する
   • count を res の末尾に追加する
3. res を返す

実装例

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

def kasai(s, suff, n):
    lcp = [0] * n
    inv = [0] * n
    for i in range(n):
        inv[suff[i]] = i
    k = 0
    for i in range(n):
        if inv[i] == n-1:
            k = 0
            continue
        j = suff[inv[i] + 1]
        while i + k < n and j + k < n and s[i + k] == s[j + k]:
            k += 1
        lcp[inv[i]] = k
        if k > 0:
            k -= 1
    return lcp

def solve(s, Q):
    res = []
    for i in range(len(Q)):
        left, right = Q[i]
        sub = s[left: right + 1]
        length = right - left + 1

        suffix = [[i, sub[i:]] for i in range(length)]

        suffix.sort(key=lambda x: x[1])
        suff, suffix = [list(t) for t in zip(*suffix)]

        lcp = kasai(sub, suff, length)
        count = len(suffix[0])
        for i in range(length - 1):
            count += len(suffix[i + 1]) - lcp[i]

        res.append(count)
    return res

s = "pptpp"
Q = [(1,1),(1,4),(1,1),(0,2)]
print(solve(s, Q))

入力

"pptpp", [(1,1),(1,4),(1,1),(0,2)]

出力

[1, 8, 1, 5]

計算量の目安

各クエリでは、接尾辞のソートに O(m log m) 回の比較(m は範囲の長さ)、Kasaiのアルゴリズムによる LCP 構築に O(m) の計算量が必要です。文字列比較のコストを考慮すると、クエリあたりの計算量は O(m² log m) 程度になります。より大規模な入力に対応したい場合は、接尾辞配列を高速に構築する手法や、Suffix Automaton(接尾辞オートマトン)など別のデータ構造の活用も検討するとよいでしょう。