Pythonで文字列とその接尾辞との類似度の合計を求めるプログラム

問題の概要

文字列 s が与えられたとき、s とそのすべての接尾辞（末尾から1文字ずつ短くした部分文字列）との「類似度」の合計を求めます。ここで2つの文字列の類似度とは、両方の文字列に共通する最長の接頭辞（先頭からの一致部分）の長さのことです。

例として、入力が s = "pqpqpp" の場合を考えてみましょう。この文字列の接尾辞は次の6つです。

• "pqpqpp"（元の文字列そのもの）
• "qpqpp"
• "pqpp"
• "qpp"
• "pp"
• "p"

それぞれと元の文字列 "pqpqpp" との類似度は、順に 6、0、3、0、1、1 となります。したがって、求める合計は 6 + 0 + 3 + 0 + 1 + 1 = 11 です。

アプローチ：Zアルゴリズムによる線形時間の解法

この問題は Zアルゴリズム の考え方を利用すると効率的に解けます。Z配列の第 i 要素 z[i] は「位置 i から始まる接尾辞と、文字列の先頭が何文字一致しているか」を表します。これはまさに、今回求めたい「接尾辞との類似度」そのものです。

各接尾辞に対して毎回先頭から素朴に比較すると計算量は O(n²) になりますが、Zアルゴリズムでは、すでに判明している一致区間 [l, r] の情報を再利用することで、文字列全体を O(n) の線形時間で処理できます。

解き方の手順

• length := 文字列 s の長さ
• total := length（最初の接尾辞は文字列自身なので、類似度は length になる。これを初期値とする）
• z := 0 を1つ含むリストとして初期化
• l := 0、r := 0（現在わかっている一致区間の左右端）
• k を 1 から length − 1 まで繰り返します。
  • k > r の場合（過去の情報が使えないため新規に比較）：
    match := 0、index := k とし、index < length のあいだ s[index] と s[match] を比較しながら両方を進めます。不一致になったらループを抜けます。
  • 求めた match を z の末尾に追加します。
  • match > 0 の場合は、total に match を加算し、l := k、r := index − 1 として一致区間を更新します。
  • k ≤ r の場合（既存の一致区間を再利用できる場合）：
    • z[k−l] < (r−k)+1 であれば、z[k−l] をそのまま z の末尾に追加し、total に加算します。
    • そうでなければ、match := r−k、index := r から比較を再開し、不一致まで進めた結果を z に追加して total に加算し、l := k、r := index − 1 を更新します。
• 最後に total を返します。

Pythonでの実装例

以下が実際の実装です。

def solve(s):
    length = len(s)
    total = length

    z = [0]
    l = 0
    r = 0

    for k in range(1, length):
        if k > r:
            # 既存の一致区間が使えないため、先頭から直接比較する
            match = 0
            index = k
            while index < length:
                if s[index] == s[match]:
                    match += 1
                    index += 1
                else:
                    break
            z.append(match)
            if match > 0:
                total += match
                l = k
                r = index - 1
        else:
            # 一致区間 [l, r] の情報を再利用する
            if z[k - l] < (r - k) + 1:
                z.append(z[k - l])
                total += z[k - l]
            else:
                match = r - k
                index = r
                while index < length:
                    if s[index] == s[match]:
                        match += 1
                        index += 1
                    else:
                        break
                z.append(match)
                total += match
                l = k
                r = index - 1
    return total

s = "pqpqpp"
print(solve(s))

入力

"pqpqpp"

出力

11

まとめ

文字列とそのすべての接尾辞との類似度の合計を求める問題は、Zアルゴリズムを使うことで O(n) の計算量で解くことができます。ポイントは、すでに判明している一致区間 [l, r] の情報を再利用して、重複する比較を省略することです。この考え方は、文字列検索やパターンマッチングなど、さまざまな文字列処理の問題にも応用できるので、ぜひ押さえておきましょう。