Pythonで列を並べ替えた後に最大の「1」だけの部分行列の面積を求めるプログラム

問題概要

0と1だけで構成される2値行列が与えられます。まず、列を好きな回数だけ自由に並べ替えることができ、その後、「1」のみで構成される最大の部分行列の面積を求めて返します。

たとえば、入力が次のような行列だった場合を考えてみましょう。

この場合の出力は 4 になります。列を並べ替えることで、次のように2行目と3行目に「1」を集めることができ、2×2の部分行列が作れるからです。

解法のアプローチ

この問題を解くには、以下の手順に従います。

• n := 行列の行数
• m := 行列の列数
• ans := 0(答えを格納する変数)
• i を 1 から n-1 まで繰り返す:
  • j を 0 から m-1 まで繰り返す:
    • matrix[i, j] が 1 の場合:
      • matrix[i, j] := matrix[i, j] + matrix[i-1, j](上の行の値を加算)
• 行列の各行に対して:
  • その行をソートする
  • j を m-1 から 0 まで 1 ずつ減らしながら繰り返す:
    • ans := ans と row[j] × (m - j) のうち大きい方
• ans を返す

考え方のポイント

まず各行について、そのセルから真上方向へ連続する「1」の個数を累積的に計算します。これにより、各行はヒストグラムの高さのように扱えるようになります。列は自由に並べ替えられるため、各行内で高さを昇順にソートすると、右端から見たときに「row[j] 以上の高さを持つ列が (m - j) 本存在する」ことになります。したがって、row[j] × (m - j) がその時点で作れる最大の長方形の面積となり、すべての場合の中での最大値が答えになります。

実装例

理解を深めるために、以下の実装を見てみましょう。

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

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

入力

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

出力

4