Pythonで行列をk個のピースに分割する方法の数を数えるプログラム

問題の概要

0と1だけで構成されるバイナリ行列と整数 k が与えられます。各ピースに必ず1つ以上の「1」が含まれるように、この行列を k 個の部分に分割する方法の数を求めるのが目的です。ただし、切断には次のルールがあり、必ずこの順序で従う必要があります。

1. 方向を選ぶ：垂直方向または水平方向のどちらかを選択します。
2. 切断位置を選ぶ：行列内のインデックスを1つ選び、そこで2つの領域に分割します。
3. 垂直に切った場合：左側の領域はこれ以上切断できず、右側だけを切り続けられます。
4. 水平に切った場合：上側の領域はこれ以上切断できず、下側だけを切り続けられます。

この条件下で、行列を分割できる異なる方法の総数を求めます。答えが非常に大きくなる可能性があるため、結果は 10^9 + 7 で割った余りを返します。

入力例と出力例

例として、次のような3×3の行列が与えられ、k = 2 の場合を考えてみましょう。

このとき出力は 4 になります。垂直方向に切れる位置が2つ、水平方向に切れる位置が2つあり、合計4通りの分割方法が存在するためです。

解法のアプローチ

この問題は、「右下からの累積和（接尾辞和）による前計算」と「再帰的な探索」を組み合わせると効率的に解けます。まず、counts[(i, j)] にセル (i, j) から右下端までの領域に含まれる「1」の個数を格納しておきます。こうすることで、任意の時点における残り領域の「1」の総数を即座に参照できるようになります。

具体的な手順は以下の通りです。

• p := 10^9 + 7（答えの剰余を取るための定数）
• m := 行列の行数、n := 行列の列数
• counts := 空の辞書（defaultdict）
• i を m − 1 から 0 まで降順にループし、さらに j を n − 1 から 0 まで降順にループしながら、counts[(i, j)] := counts[(i + 1, j)] + counts[(i, j + 1)] − counts[(i + 1, j + 1)] + matrix[i][j] を計算する
• 関数 f(x, y, c) を定義する（左上隅が (x, y) の領域に対して、残り c 回の切断を行う方法の数を返す）
  • count := counts[(x, y)]（現在の領域に含まれる「1」の総数）
  • c が 0 の場合：count > 0 なら 1、そうでなければ 0 を返す
  • ans := 0
  • i を x + 1 から m − 1 までループ：0 < counts[(i, y)] < count を満たすなら、ans := ans + f(i, y, c − 1）（水平方向の切断）
  • j を y + 1 から n − 1 までループ：0 < counts[(x, j)] < count を満たすなら、ans := ans + f(x, j, c − 1）（垂直方向の切断）
  • ans mod p を返す
• 最後に f(0, 0, k − 1) を呼び出した結果を返す

切断位置の候補では、切り離される側（左または上）にも「1」が含まれている必要があります。counts[(i, y)] は切断後に残る下側領域の「1」の個数を表すため、これが 0 より大きく、かつ現在の総数 count より小さい（＝上側にも「1」が残る）ことが条件になります。

Pythonでの実装例

それでは、理解を深めるために実際のコードを見てみましょう。

from collections import defaultdict


class Solution:
    def solve(self, matrix, k):
        p = 10 ** 9 + 7

        m, n = len(matrix), len(matrix[0])
        counts = defaultdict(int)
        for i in range(m)[::-1]:
            for j in range(n)[::-1]:
                counts[(i, j)] = (
                    counts[(i + 1, j)]
                    + counts[(i, j + 1)]
                    - counts[(i + 1, j + 1)]
                    + matrix[i][j]
                )

        def f(x, y, c):
            count = counts[(x, y)]
            if c == 0:
                return 1 if count > 0 else 0

            ans = 0
            for i in range(x + 1, m):
                if 0 < counts[(i, y)] < count:
                    ans += f(i, y, c - 1)
            for j in range(y + 1, n):
                if 0 < counts[(x, j)] < count:
                    ans += f(x, j, c - 1)

            return ans % p

        return f(0, 0, k - 1)


ob = Solution()
matrix = [
    [1, 1, 0],
    [1, 0, 1],
    [1, 1, 1],
]
k = 2
print(ob.solve(matrix, k))

入力

[
    [1, 1, 0],
    [1, 0, 1],
    [1, 1, 1]
], 2

出力

4

まとめ

本記事では、バイナリ行列を「1」を1個以上含む k 個のピースに分割する方法の数を求める問題を解説しました。右下からの累積和で各領域の「1」の個数を前計算しておくことで、切断ごとの妥当性チェックを高速に行えます。探索の状態数は O(m × n × k)、各状態の遷移は O(m + n) となるため、行列サイズが大きいケースでは functools.lru_cache などによるメモ化を併用すると、さらなる高速化が期待できます。