問題1.7

世界で闘うプログラミング力を鍛える150問

寝る前に一問。

問題

MxNの行列について、要素が0であれば、その行と列のすべてを0にするようなアルゴリズムを書いてください。

#encoding: utf-8


def initialize_with_zero(matrix):
    '''
    引数で受け取ったMxNの行列で、要素が0であれば、その行と列の全てを0にする関数
    >>> initialize_with_zero([[1, 1, 1], [1, 0, 1], [1, 1, 1]])
    [[1, 0, 1], [0, 0, 0], [1, 0, 1]]
    >>> matrix = [[0, 1, 1, 0], [1, 1, 1, 1], [1, 1, 1, 1], [0, 1, 1, 0]]
    >>> initialize_with_zero(matrix)
    [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]]
    '''
    # 要素が0である箇所を記憶するリスト
    zero_list = []

    # 要素が0である箇所を探す
    for row_counter in range(0, len(matrix)):
        for column_counter in range(0, len(matrix[0])):
            if matrix[row_counter][column_counter] == 0:
                zero_list.append({'row_counter': row_counter,
                                  'column_counter': column_counter})

    # 要素が0である箇所の列と行を全て0にする
    for zero_target in zero_list:
        for row_counter in range(0, len(matrix)):
            matrix[row_counter][zero_target['column_counter']] = 0
        for column_counter in range(0, len(matrix[0])):
            matrix[zero_target['row_counter']][column_counter] = 0

    return matrix


def _test():
    import doctest
    doctest.testmod()


if __name__ == '__main__':
    _test()

備考

計算量的によくない解き方だったようです。計算量がO(MN)となる。
要素が0の箇所を行と列で記憶するんじゃなくて、要素が0を含む行、要素が0を含む列で記憶しておけば、
0をセットする時にもっとシンプルになりますね。。。