【備忘録】競プロ典型90問 004(★2) #2

AtCoderで公開されている、E869120さんの『競プロ典型90問』を解いていきます。

---

　前回の記事に引き続き、004 - Cross Sumをやっていきます。今回はリスト内包表記、zip関数、lambda式といった、コードを簡潔に書くテクニックについて学んでいきます。

現状の課題

　まず、こちらが前回の記事にあげた僕のコードです。forループを8回使っており、何をしているのかがよく分かる反面、コードが長ったらしいです。

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

# ここで各行の和rowsumと各列の和colsumを計算しておく
rowsum = []
for i in range(H):
   tmp = 0
   for j in range(W):
       tmp += A[i][j]
   rowsum.append(tmp)  
colsum = []
for j in range(W):
   tmp = 0
   for i in range(H):
       tmp += A[i][j]
   colsum.append(tmp)
   
B = [[0] * W for i in range(H)]
for i in range(H):
   for j in range(W):
       B[i][j] = rowsum[i] + colsum[j] - A[i][j]

for b in B:
   print(*b)

　こちらはdrkenさんによる解答例になります。forループは何とたったの2回で、かなり簡潔に書かれていることが分かります。

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

# 前処理
yoko = list(map(sum, A))
tate = list(map(sum, zip(*A)))

# 各マス
for i in range(H):
   print(' '.join(map(lambda j: str(yoko[i] + tate[j] - A[i][j]), range(W)))

出典：

　コードを簡潔に記述するためのテクニックとして、

• リスト内包表記
• map関数
• zip関数
• lambda式
• わざわざ新しいリストを作らない

などが挙げられそうです。既に使っているものもありますが、使いこなせるようになるために、改めて調べてアウトプットをしていきます。

リスト内包表記

　リスト内包表記は、既存のリストから新しいリストを作る際に有用であるという認識でした。例えば

B = []
for a in A:
   a = a + 1
   B.append(a)

# 等価

B = [a + 1 for a in A]

のように書けます。つまり、appendを省略することができて、一般化すれば

new_lst = []
for old_element in old_lst:
    new_element = function(old_element)
    new_lst.append(new_element)

# 等価

new_lst = [function(old_element) for old_element in old_lst]

この2つが等価ということになり、リスト内包表記を用いればかなりコードを短く書くことができます。なお、今回の問題のように入力から行列を二次元配列で受け取るような場合にも同様に

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

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

のように書けます。今後はいちいちappendしないように気を付けます。

map関数

　map関数は入力を受け取る際に便利という認識でした。

map(function, iterable)

の形で用いられ、第二引数のiterableに順番に関数functionを作用させたものを返します。返り値には気を付ける必要があり、主に

# iterableの長さが予め決まっていて十分短い場合
A_1, A_2, ... A_i, ..., A_N = map(function, iterable) # N = len(iterable)

# list()を用いてlist型に変換する
A = list(map(function, iterable))

この2つの受け取り方があります。

　なお、今回の問題では入力時に限らず、各行・各列の和を計算する際に

rowsum = list(map(sum, A))

のように用いることでかなり簡潔に書くことができます。この例では、二次元配列Aの各要素（一次元配列たち）のそれぞれに、iterableの全要素の和を求めるsum関数を作用させることで、行列Aの各行の和を求め、リストにしてrowsumに引き渡しています。

zip関数

　zip関数は複数のイテラブルを引数に受け取り、1つずつ要素を取り出してまとめてくれる関数のようです。なんて便利なんだ、笑。なお、出力はmap関数のように癖があるようで、list()でリスト化するのがよさそうです。

list(zip(*A)) # = zip(A[0], A[1], A[2], ..., A[H-1])

のように用いることで、二次元配列$A$から、列方向の一次元配列を$W$個まとめた二次元配列を得ることができます。

　解答例のように、和を求めるだけでよいならリスト化せずにそのままmap関数に渡すこともできるという、、、便利ですね。勉強になります。

lambda式（無名関数）

　わざわざ新しい関数を作るほどでもないなぁ、っていうときに使えるらしいです。

# 定義して
def func(引数):
   return 返り値
# 値を代入して使う
func(値)

# 等価
# 定義、代入を一度にしてしまう
(lambda 引数: 返り値)(値)

再利用の可能性がないときには積極的にlambda式を使うことで短くコードを書くことができそうです。解答例のようなmap関数との合わせ技も定番のようです。

まとめ

　今までforループのごり押しで解決しようとしていた節がありましたが、必ずと言っていいほどコードを短く書く工夫は存在するのでしょう。今後はもっともっと調べて、どんどん吸収していこうと思います。今後の僕の回答例で成長した姿を見せることができればと思います！

　お読みいただきありがとうございました。

参考記事