Pythonで数値文字列を分割して値のリストを作れるパターン数をカウントするプログラム

問題概要

0〜9の数字のみから構成される文字列 s と、整数 k が与えられます。このとき、s を [1, k] の範囲に含まれる数値からなるリストとして分割できる方法が何通りあるかを求めます。答えが非常に大きくなる可能性があるため、結果は 10^9 + 7 で割った余りを返します。

たとえば、入力が s = "3456"、k = 500 の場合、出力は 7 になります。これは次の7通りの分割が可能であるためです。

• [3, 4, 5, 6]
• [34, 5, 6]
• [3, 4, 56]
• [3, 45, 6]
• [34, 56]
• [345, 6]
• [3, 456]

解法の考え方：動的計画法（DP）

この問題は動的計画法を用いることで効率よく解けます。dp[i] を「部分文字列 s[i:] を条件を満たす数値リストに分割する方法の総数」と定義します。文字列の末尾側から順に計算していくことで、各位置からの分割パターン数を累積的に求められます。

具体的には、位置 i から始まる部分文字列に対して、区切り位置 j を順に試します。s[i..j] を数値として読み取り、その値が 1 以上 k 以下であれば、その区切りを採用した場合の残りのパターン数 dp[j + 1] を dp[i] に加算します。一度 curr_val が k を超えたら、それ以上桁を伸ばしても値は増える一方なので、内側のループを抜けて次へ進みます。

アルゴリズムの手順

• m := 10^9 + 7

• N := 文字列 s の長さ

• dp := サイズ (N + 1) のリストを作成し、すべて 0 で初期化する

• dp[N] := 1（残りの文字列が空の場合の分割方法は1通りとみなす）

• i を N − 1 から 0 まで 1 ずつ減らしながら繰り返す：

  • curr_val := 0

  • j を i から N − 1 まで 1 ずつ増やしながら繰り返す：

    • curr_val := curr_val × 10 + （s[j] を数値に変換した値）

    • もし curr_val が 1 以上 k 以下であれば：

      • dp[i] := (dp[i] + dp[j + 1]) mod m

    • そうでなければ：

      • 内側のループを抜ける

• dp[0] を答えとして返す

実装例

以下は、上記のアルゴリズムをPythonで実装したコードです。

class Solution:
    def solve(self, s, k):
        m = 10 ** 9 + 7
        N = len(s)
        dp = [0] * (N + 1)
        dp[N] = 1
        for i in range(N − 1, −1, −1):
            curr_val = 0
            for j in range(i, N):
                curr_val = curr_val * 10 + int(s[j])
                if 1 <= curr_val <= k:
                    dp[i] = (dp[i] + dp[j + 1]) % m
                else:
                    break
        return dp[0]
ob = Solution()
s = "3456"
k = 500
print(ob.solve(s, k))

入力

"3456", 500

出力

7

計算量の目安

内側のループは curr_val が k を超えた時点で打ち切られるため、各開始位置 i について調べる桁数は高々 k の桁数（d 桁）にとどまります。したがって、時間計算量は O(N × d)、空間計算量は dp 配列ぶんの O(N) となります。全探索に比べて大幅に効率的であり、文字列が長い場合でも現実的な時間で答えを求められます。