Pythonで二次元リストを簡潔にコーディングする(リスト内包表記)

この記事はPythonの二次元配列を簡潔にコーディングする方法について記載する．

標準入力

標準入力が以下のようだったとする

Pythonでの内包表記

内包表記（List Comprehension）を使って，ユーザーから入力された行列 A を二次元リストとして作成する．

example.py

H, W = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(H)]

map()は配列の要素を第一引数に指定された型に変換する
map(int, ["1", "2", "3"])ならば[1, 2, 3]となる．`

for _ in range(N)はN 回のループを実行する．_ は，ループ変数が使用されないことを示すために一般的に使用される変数名である．

for _ in range(H):
H 回のループを実行する．
_ は,ループ変数が使用されないことを示すために一般的に使用される変数名．

input().split()
ユーザーから入力された文字列を空白で分割し，文字列のリストを返す．

競プロの標準入力ではよく使うので覚えておこう

map(int, ...)
map() 関数を使って，分割された文字列のリストの各要素を整数に変換する．

list(...)
map() 関数の結果をリストに変換する．

最終的に，A は H 行 W 列の整数からなる二次元リストになる．

全ての要素が0の二次元リストを作る

H, W = map(int, input().split())
B = [[0] * W for _ in range(H)]

これで全ての要素が0であるH×Wの二次元リストができる．

二次元リストに出力を[]を使わずに

for row in B:
    print(row)

だと，以下のような出力になる時，

[5, 5, 5]
[5, 5, 5]
[5, 5, 5]

[]を使わずに以下のように出力するには，

5 5 5
5 5 5
5 5 5

以下のようなコードにすればいい．

example.py

for row in B:
    print(*row)

B の各行 row について，*を使って要素をアンパックし，空白区切りで出力する．

二次元リストの処理でTLEを防ぐには

各i行j列における和を全て計算して結果を二次元リストに格納したい時
ある行(または列)での全ての要素の和を計算する時，各i行j列それぞれでsum関数を使うと非常に計算量が多いコードとなる．

bad_example.py

for i in range(H):
    row_sum = sum(A[i])
    for j in range(W):
        col_sum = sum(A[k][j] for k in range(H))
        B[i][j] = row_sum + col_sum - A[i][j]

各マス (i, j) について，行の合計と列の合計を計算するために，それぞれ sum(A[i]) と sum(A[k][j] for k in range(H)) を使用している．これらの計算は，各マスごとに行われるため，時間計算量が O(H^2 * W) になる．与えられた制約のH と W の最大値が 大きい場合，このアプローチでは大きなケースでTLEになる可能性が高い．そのため以下のようなアプローチを取ると良い．

good_example.py

# 行の合計を計算
row_sums = [sum(row) for row in A]

# 列の合計を計算
col_sums = [sum(A[i][j] for i in range(H)) for j in range(W)]

# 各マスの値を計算(A[i][j]は行と列の合計計算で重複しているので減算)
B = [[row_sums[i] + col_sums[j] - A[i][j] for j in range(W)] for i in range(H)]

行の合計を計算では何が行われているのか
for row in A:
行列 A の各行 row に対してループを実行する．

sum(row)
各行 row の要素の合計を計算する．

[... for row in A]
内包表記を使って，各行の合計を要素とするリストを作成．

列の合計を計算では何が行われているのか
for j in range(W):
列のインデックス j に対してループを実行する．(W は列数)

sum(A[i][j] for i in range(H))
各列 j の要素の合計を計算する．内包表記を使って，行のインデックス i に対してループを実行し，A[i][j] の値を合計する．(H は行数)

[... for j in range(W)]
内包表記を使って，各列の合計を要素とするリストを作成する．

前処理で行の合計と列の合計を計算しておくことで，各マスの値を効率的に計算できるようになる．
このアプローチにより，時間計算量が O(H * W) に削減され，TLEを回避できる．