Pythonで最大k回の連続勝利となるゲーム結果の組み合わせ数を数える方法

問題の概要

2つの整数 n と k が与えられます。n はこれから行う試合の総数、k は許容される連続勝利の上限です。全 n 試合の勝敗(勝ち=W、負け=L)の並びのうち、「連続した勝利が k 回以下」に収まるパターンが何通りあるかを求めます。答えは非常に大きくなる可能性があるため、10⁹ + 7 で割った余りを返します。

例えば n = 3、k = 2 の場合、答えは 7 です。条件を満たすパターンは次の 7 通りです。

LLL / WLL / LWL / LLW / WWL / LWW / WLW

解き方のアプローチ

この問題は再帰的な動的計画法(DP)で解くことができます。i 番目の試合を処理する局面で、現在の連続勝利数を K とした関数 dp(i, K) を定義します。

• i 試合目で負ける場合は、連続勝利数がリセットされるため dp(i + 1, 0) に遷移します。
• i 試合目で勝つ場合は、連続勝利数が1増えるため dp(i + 1, K + 1) に遷移します。
• すべての試合を消化したとき(i ≥ n)、または連続勝利が k を超えた時点で再帰を打ち切り、その経路が条件を満たしていれば 1、満たしていなければ 0 を返します。
• 最終的に dp(0, 0) の値を 10⁹ + 7 で割った余りが答えになります。

Pythonでの実装例

def solve(n, k):
    m = 10**9 + 7

    def dp(i, K):
        # 全試合を消化、または連続勝利がkを超えたら判定
        if i >= n or K > k:
            return 1 if K <= k else 0
        # 負ける場合と勝つ場合を足し合わせる
        return (dp(i + 1, 0) + dp(i + 1, K + 1)) % m

    return dp(0, 0)

n = 4
k = 2
print(solve(n, k))

入力

n = 4, k = 2

出力

13

n = 4、k = 2 の場合、全 16 通りのうち WWWL・LWWW・WWWW の 3 パターンだけが「3回以上の連続勝利」に該当するため、残りの 13 通りが答えとなります。

実装上の注意点

• 剰余の法は必ず 10**9 + 7(10の9乗+7)と記述してください。誤って 1**9 + 7 と書くと 1 + 7 = 8 として計算されてしまい、まったく異なる結果になります。
• 上記の素朴な再帰は同じ状態を繰り返し計算するため、n が大きくなると指数時間がかかります。次に紹介するメモ化で大幅に高速化できます。

メモ化による高速化

標準ライブラリの functools.lru_cache を使えば、一度計算した状態をキャッシュでき、計算量を O(n × k) まで抑えられます。

from functools import lru_cache

def solve(n, k):
    m = 10**9 + 7

    @lru_cache(maxsize=None)
    def dp(i, K):
        if i >= n or K > k:
            return 1 if K <= k else 0
        return (dp(i + 1, 0) + dp(i + 1, K + 1)) % m

    return dp(0, 0)

print(solve(100, 5))

まとめ

連続勝利の制約付き組み合わせ数の問題は、「現在の試合番号」と「現在の連続勝利数」を状態とする再帰DPで自然に表現できます。メモ化を併用すれば大きな n・k に対しても効率的に求解できるため、競技プログラミングでも応用範囲の広いテクニックです。