Pythonで連続文字の制限を満たす同サイズの文字列を数えるプログラム

小文字の英字だけで構成された文字列 s と整数 k が与えられます。このとき、次の3つの条件をすべて満たす文字列の総数を求めるのが今回の課題です。

• s と同じ長さである
• 辞書順で s 以下である
• 同じ文字が連続する回数が k 以下である(k を超える連続は不可)

答えは非常に大きくなる可能性があるため、10^9 + 7 で割った余りとして返します。

たとえば入力が s = "app"、k = 2 のとき、条件を満たす文字列は 405 個あるため、出力は 405 になります。

解法の考え方:桁DP

「指定した文字列以下の文字列を数え上げる」タイプの問題では、桁DP(デジットDP)と呼ばれる手法が定番です。文字列の先頭から1文字ずつ決めながら、次の状態を管理して再帰的に探索します。

• pos:現在決めている文字の位置
• bound:ここまでの文字が s と完全一致しているか(一致中なら次の文字は s の対応する文字までしか選べない)
• last:直前に置いた文字
• count:last が連続している回数

処理の手順をまとめると次のようになります。

• k <= 0 なら、条件を満たす文字列は存在しないので 0 を返す。
• 剰余用の定数 m = 10^9 + 7 を用意する。
• n に文字列 s の長さを代入する。
• nums に、s の各文字を ord(char) - ord("a") で 0〜25 の数値に変換したリストを作る。
• dp(0, True, -1, 0) の結果を m で割った余りを返す。

dp 関数の中身

dp 関数は引数として pos(現在位置)、bound(上限に達しているか)、last(直前の文字)、count(連続回数)を受け取り、次のように動作します。

• count > k の場合:連続制限を超えているので 0 を返す。
• pos == n の場合:最後まで文字列を組み立てられたので 1 を返す。
• num = nums[pos]:今の位置で選べる文字の上限値。
• res = 0 から始め、bound が真なら 0〜num、偽なら 0〜25 の範囲で各文字 i を試す。
• 各 i について res += dp(pos + 1, bound and i == num, i, count * (i == last) + 1) を加算する。
  ※ i == last が真なら count + 1、偽なら 1 になるため、連続カウントが正しく更新されます。
• 最終的に res を返す。

Pythonでの実装例

class Solution:
    def solve(self, s, k):
        if k <= 0:
            return 0
        MOD = 10 ** 9 + 7
        n = len(s)
        nums = [ord(char) - ord("a") for char in s]

        def dp(pos, bound, last, count):
            if count > k:
                return 0
            if pos == n:
                return 1
            num = nums[pos]
            res = 0
            for i in range(num + 1 if bound else 26):
                res += dp(pos + 1, bound and i == num, i, count * (i == last) + 1)
            return res

        return dp(0, True, -1, 0) % MOD


ob = Solution()
print(ob.solve('app', 2))

入力

s = "app"
k = 2

出力

405

動作のポイント

この実装では、最初の文字は必ず "a"(s の先頭文字)しか選べず、以降も bound が立っている間は s の該当文字までしか選択できません。一度 s より小さい文字を選んだ時点で bound が外れ、残りの位置は 26 種類の文字から自由に選べるようになります。

なお、上記のコードはシンプルな素朴再帰のため、文字列が長くなると計算量が急増します。実務では functools.lru_cache などを利用して (pos, bound, last, count) をキーにメモ化すると、状態数は高々 n × 2 × 26 × (k+1) に抑えられるため、多項式時間で効率的に処理できます。