【Python】行列内の0を含む行と列をすべて0に変換するアルゴリズム

問題概要

2次元の数値行列が与えられたとき、行列内にある0を含む行と列のすべての要素を0に置き換え、最終的な行列を返すプログラムを作成します。

例えば、入力行列の中で0が存在する行(0行目、2行目、3行目)は、最終的な行列ではすべて0になります。同様に、0が存在する列(0列目、1列目、2列目)もすべて0に変換されます。

解法のアプローチ

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

1. 行数を n、列数を m とします。
2. n × m のサイズの結果用行列 res を作成し、すべての要素を0で初期化します。
3. 元の行列の転置行列 transpose を作成します。
4. 各行 i について、matrix[i] に0が含まれていない場合、各列 j に対して transpose[j] にも0が含まれていないかを確認します。両方に0が存在しなければ、res[i][j] に元の値 matrix[i][j] を代入します。
5. 最後に res を返します。

擬似コードで表すと次のようになります。

n := 行数, m := 列数
res := n x m のサイズの行列を作成し、0で埋める
transpose := 与えられた行列の転置行列
for each row i, do
    if 0 not in matrix[i], then
        for each column j in matrix, do
            if 0 not in transpose[j], then
                res[i, j] := matrix[i, j]
return res

Pythonでの実装例

それでは、実際のPythonコードを見てみましょう。

class Solution:
    def solve(self, matrix):
        n, m = len(matrix), len(matrix[0])
        res = [[0 for __ in range(m)] for _ in range(n)]
        transpose = [list(row) for row in zip(*matrix)]

        for i in range(n):
            if 0 not in matrix[i]:
                for j in range(m):
                    if 0 not in transpose[j]:
                        res[i][j] = matrix[i][j]

        return res

ob = Solution()
matrix = [
    [6, 0, 0, 6, 9],
    [4, 9, 9, 4, 8],
    [0, 8, 3, 4, 2],
    [9, 0, 7, 8, 3],
    [5, 2, 9, 6, 8]
]
print(ob.solve(matrix))

入力

matrix = [
    [6, 0, 0, 6, 9],
    [4, 9, 9, 4, 8],
    [0, 8, 3, 4, 2],
    [9, 0, 7, 8, 3],
    [5, 2, 9, 6, 8]
]

出力

[[0, 0, 0, 0, 0], [0, 0, 0, 4, 8], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 6, 8]]

コードのポイント

このアルゴリズムの鍵となるのは転置行列の活用です。zip(*matrix) を使って転置行列を作成することで、「列 j に0が含まれているか」という判定を、リストの in 演算子でシンプルに記述できるようになります。

計算量は O(n × m × (n + m)) となります。これは、各セル (i, j) ごとに行と列の走査が必要になるためです。もしパフォーマンスをさらに向上させたい場合は、事前に0を含む行番号と列番号をセット(set)に記録しておく方式に変更することで、計算量を O(n × m) まで削減できます。