はじめに
こんにちは！ 中2のAwashAmityOakです！ AtCoderでは、早く解くために、解法を思いつく思考力だけでなく、コードを短く読みやすく書くことも重要です。今回は、私が実践している、リストを受け取るときのTipsを広める記事です。
良くないコード
例えば、次のような形式で、入力が与えられるとしましょう。すべて数値とします。
$N \quad M$
$A_1, \quad A_2, \quad \dots, \quad A_N$
$B_1, \quad B_2, \quad \dots, \quad B_M$
このとき、Pythonでは、
N, M = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
と書くことができます。
しかし、map()の前後がlist()で囲まれていて、少しタイプしにくいかと思います。
改善点
そこで、私が提案するリストの受け取り方はこんな書き方です。
N, M = map(int, input().split())
*A, = map(int, input().split())
*B, = map(int, input().split())
どうですか？ きれいに見えませんか？
Aの部分で*A,という書き方をすることで、list()がいらず、3行ともすべてmap(int, input().split())という統一された書き方ができます。
何をしているか
Pythonのアンパック(Unpack)という構文を使っています。(「分割代入」と呼ばれることもあります。)
上記のコードのNとMのように、変数の個数(この場合は2個)が決まっている場合は、N, M = ...のように、まとめて代入式を書くことができます。この書き方をアンパックと言います。アンパックでは、代入式の右辺はイテレータである必要があります。(map()の戻り値は、map型というイテレータです。)これは、競プロの入力の受け取り方などを調べると、たくさん情報が出てくるので、知っている人も多いかと思います。
しかし、アンパックがすごいのは、これだけではありません。*(アスタリスク)を使うことで、変数の個数が決まっていなくても、「残り全部」をリストとして代入することができるのです。
a, b, *c = [0, 1, 2, 3, 4, 5]
例えば、このコードでは、aに0、bに1が代入されるのは、一目瞭然ですが、*がついているcには、aとbに入った0と1以外のすべての要素([2, 3, 4, 5])がリストとして代入されます。
*を使う位置は、変えることも可能です。
a, *b, c = [0, 1, 2, 3, 4, 5]
この場合は、aに0、bに[1, 2, 3, 4]、cに5が代入されます。
*をつける変数の数は、0個か1個でなければなりません。
# エラーになる
a, *b, *c = [0, 1, 2, 3, 4, 5]
変数の数が1個で、その変数に*でアンパックで代入したい場合は、次のコードのように、変数名の後ろに,(コンマ)をつける必要があります。
*A, = map(int, input().split())
先ほどのコードの一部です。(分割代入とか言いながら、もはや分割していませんけどね…。)
応用
クエリ問題で、以下の形式のクエリを受けることがあります。(先日のABC344のE問題がこれでした。)
次の2種類のどちらかがクエリとして渡される。
1 x y
2 x
このときは、入力を受け取るまで、変数の個数が3個と2個のどちらなのかわからないですよね？ このようなとき、アンパックを駆使して、こんな書き方をすると、スマートです。
# ...前略
for q in range(Q):
t, *a = map(int, input().split())
if t == 1:
x, = a
# 処理...
elif t == 2:
x, y = a
# 処理...
# 後略...
(実際のコンテスト中は、きれいなコードを書くことを、意識する必要なんてまったくありませんが…。)
おわりに
今回は、アンパックを使用して、AtCoderのコードを短く綺麗に書く方法をお伝えしていきました。
参考
この記事が良いと思ったら、ぜひいいねとフォローをよろしくお願いします！ Xアカウントのフォローなどもしてくれると泣いて喜びます！
最後まで読んでくださり、ありがとうございました！