Pythonで列を並べ替えた後に最大の部分行列を見つける方法

問題概要

m×n のバイナリ行列（各要素が 0 か 1 のみで構成される行列）が与えられます。この行列の列は、任意の順序で自由に入れ替えることができます。列の並べ替えを行った後、すべての要素が 1 であるような最大の部分行列を見つけ、その面積を求めるのがこの問題の目的です。

たとえば、入力が次のような行列だったとしましょう。

このとき出力は 4 になります。列を入れ替えることで、次のような行列が得られるからです。

オレンジ色で示した部分が最大の部分行列で、1 が 4 つ並ぶ 2×2 の正方形になっています。面積は 4 です。

解法のステップ

この問題は、次の手順で解くことができます。

• row := 行列の行数、col := 行列の列数 とする
• j を 0 から col−1 まで繰り返す
  • i を 1 から row−1 まで繰り返す
    • matrix[i][j] が 1 ならば、matrix[i][j] := matrix[i][j] + matrix[i−1][j] と更新する
• ans := 0 とする
• i を 0 から row−1 まで繰り返す
  • リスト matrix[i] を昇順にソートする
  • j を col−1 から 0 まで 1 ずつ減らしながら繰り返す
    • matrix[i][j] が 0 ならば、ループを抜ける
    • ans = max(ans, (col−j) × matrix[i][j]) と更新する
• ans を返す

なぜこのアルゴリズムでうまくいくのか

最初の二重ループでは、各マスについて「そのマスで終わる、縦方向に連続する 1 の個数（高さ）」を累積的に計算しています。イメージとしては、ヒストグラムの棒の高さを求める処理に近いものです。

続いて、各行ごとに高さのリストをソートします。列は自由に並べ替えられるため、高い列ほど右側に集めればよいことになります。ソート後のリストで右端から k 個の列を選ぶと、その範囲をすべて覆う長方形の高さは「選んだ k 列の中で最小の高さ」、つまり matrix[i][col−k] になります。したがって、面積は (col−j) × matrix[i][j] として計算でき、すべての行と幅の組み合わせについて最大値を取れば答えが得られます。

実装例

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

def solve(matrix):
    row, col = len(matrix), len(matrix[0])
    for j in range(col):
        for i in range(1, row):
            if matrix[i][j]:
                matrix[i][j] += matrix[i-1][j]
    ans = 0
    for i in range(row):
        matrix[i].sort()
        for j in range(col-1, -1, -1):
            if matrix[i][j] == 0:
                break
            ans = max(ans, (col-j)*matrix[i][j])
    return ans

matrix = [[0,0,1],[1,1,1],[1,0,1]]
print(solve(matrix))

入力

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

出力

4

計算量

時間計算量は O(row × col × log col) です。各行のソートに O(col log col) かかり、それを row 回繰り返すためです。空間計算量については、入力行列をそのまま書き換えて利用しているため、追加で O(1) しか必要としません。