Pythonで同種（ホモジニアス）な部分文字列の数をカウントするプログラム

文字列 s が与えられたとき、その中に含まれる「同種（homogenous）な部分文字列」の総数を求める問題を考えます。同種な文字列とは、構成するすべての文字が同一である文字列のことです。答えは非常に大きな値になる可能性があるため、10^9+7 で割った余りを返します。

問題の例

たとえば、入力が s = "xyyzzzxx" の場合、出力は 13 になります。同種の部分文字列は次のように数えられます。

• "x" が 3 回

• "xx" が 1 回

• "y" が 2 回

• "yy" が 1 回

• "z" が 3 回

• "zz" が 2 回

• "zzz" が 1 回

したがって、(3 + 1 + 2 + 1 + 3 + 2 + 1) = 13 となります。

解法のアプローチ

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

1. s の末尾に番兵（センチネル）として "@" を連結します。これにより、ループ内で最後の連続区間も確実に締めくくられるようになります。

2. 空の辞書 h を用意し、prev := s[0]、c := 1 と初期化します。

3. s のインデックス 1 以降の各文字 i について処理を行います。

   • prev と i が異なる場合：これまでの連続文字列 prev * c をキーとして辞書 h に出現回数を記録し、カウント c を 1 にリセットします。

   • prev と i が同じ場合：カウント c を 1 増やします。

   • 最後に prev := i と更新します。

4. 合計値 fin := 0 とします。

5. 辞書 h の各キー（連続文字列）について、長さ t の連続区間から作れる部分文字列の総数 k = t + (t-1) + … + 1 を計算し、fin += k * h[i] として加算します。

6. fin mod (10^9+7) を返します。

ここでポイントとなるのは、長さ t の同一文字の連続区間からは、t + (t-1) + … + 1 = t(t+1)/2 個の同種部分文字列が作れるという点です。この性質（三角数）を使うことで、各区間ごとに効率よく個数を集計できます。

実装例

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

def solve(s):
    s += "@"
    h = {}
    prev = s[0]
    c = 1
    for i in s[1:]:
        if prev != i:
            if prev * c in h:
                h[prev * c] += 1
            else:
                h[prev * c] = 1
            c = 1
        else:
            c += 1
        prev = i

    fin = 0
    for key in h:
        t = len(key)
        k = 0
        while t != 0:
            k += t
            t -= 1
        fin += k * h[key]
    return fin % 1000000007

s = "xyyzzzxx"
print(solve(s))

入力

"xyyzzzxx"

出力

13

補足：よりシンプルな最適化版

辞書を使わず、連続区間ごとに三角数の公式 t(t+1)/2 を直接適用すれば、より簡潔に書けます。

def solve(s):
    MOD = 10**9 + 7
    total = 0
    run = 1
    for i in range(1, len(s)):
        if s[i] == s[i - 1]:
            run += 1
        else:
            total += run * (run + 1) // 2
            run = 1
    total += run * (run + 1) // 2
    return total % MOD

print(solve("xyyzzzxx"))  # 13

どちらの方法でも計算量は O(n) となり、長い文字列に対しても高速に動作します。