stMind

about Tech, Computer vision and Machine learning

Pythonで行列から行と列を取得するコードを簡潔に記述する

tl;dr

アスタリスクによるアンパックとzipを組み合わせる。

for row_vals, col_vals in zip(matrix, zip(*matrix)):
    print(row_vals, col_vals)

簡単な説明

まず、matrixは次のようなリストとする。

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

行は難しいことはなく、問題は列の方。 ここで、アスタリスクでアンパック→アンパックされた各リストに対して、zipで一つずつ要素を取り出してタプルとして返す、という処理をすることで行列の転置と同じことが実現できる。

print(*matrix)
# [1, 2, 3] [4, 5, 6] [7, 8, 9]

for col_vals in zip(*matrix):
    print(col_vals)
# (1, 4, 7)
# (2, 5, 8)
# (3, 6, 9)

従って、最初のコードを実行すると、行と列を一つずつ取り出すことが出来る。

for row_vals, col_vals in zip(matrix, zip(*matrix)):
    print(row_vals, col_vals)
# [1, 2, 3] (1, 4, 7)
# [4, 5, 6] (2, 5, 8)
# [7, 8, 9] (3, 6, 9)