競技プログラミングにおける、標準入力の取得例 in Python3

問題を解く以前に、入力をどう取り扱ったらいいかイマイチ良く分からない、あるいは入力を取得するコードを書くところで悩んで時間を食ってしまう……という人向けの記事。

慣れれば誰でも簡単に書けるようになるので、一応それらしいパターンをまとめつつ、どうしてこういう書き方になるのか、ちょっとした解説をつけてみようと思う。

最初に一連のコード例を。

# A: 一行の文字列を受け取る
# abcde
string = input()

# B: 一行の整数を受け取る
# 12345
integer = int(input())

# C: 一行の小数を受け取る
# 1.2345
float_number = float(input())

# D: 最初の行に要素数Nが指定され、次行のスペースで区切られた整数列を受け取る
# 5
# 1 2 3 4 5
N = int(input())
integer_list = list(map(int, input().split()))

# D': スペースで区切られた2つの整数N, Mを受け取る
# 3 4
N, M = map(int, input().spilt())

# E: 最初に要素数Nが指定され、続くN行の整数を受け取る
# 3
# 1
# 2
# 3
N = int(input())
integer_list = [int(input()) for _ in range(N)]

# F: 最初に要素数N、行数Mが指定され、続くM行の整数列（要素数N）を受け取る
# 2 3
# 1 2
# 3 4
# 5 6
N, M = map(int, input().split())
integer_matrix = [list(map(int, input().split())) for _ in range(M)]

Pythonでは、とりあえずinput()を使用することで標準入力を「文字列として」一行分受け取ることができる。実に簡単だ。

文字列ならばそのまま受け取ればいいし、整数ならint(input())、小数ならfloat(input())と型変換をしてやればいい。（例A～C）

# D: 最初の行に要素数Nが指定され、次行のスペースで区切られた整数列を受け取る
# 5
# 1 2 3 4 5
N = int(input())
integer_list = list(map(int, input().split()))

例Dは、細かく書いていくと次のような処理を行っている。一行目で要素数を受け取りはするが、この中では使用されない。

N = int(input())
input_string = input() # "1 2 3 4 5"
splitted_string = input_string.split() # ["1", "2", "3", "4", "5"]
map_object = map(int, splitted_string) # map_object(1, 2, 3, 4, 5)
integer_list = list(map_object) # [1, 2, 3, 4, 5]

mapは関数とリストを引数に取る少し特殊な関数で、大まかにいうと次のような感じの処理を行ってくれる。

# map(function, list)とした場合
result = []
for item in list:
    result.append(function(item))
return result

今回の例でいうと、リストの各要素（文字列）にint()を適用し、整数に変換してくれたものを取得できるということだ。[a, b, c]→[int(a), int(b), int(c)]といったイメージである。

ただし、map()の返り値は実はリストそのものではなく、map objectという繰り返し処理可能（イテラブル）な特殊オブジェクトのため、そのままでは扱いにくい（特に、インデックスが使えない）。そこで、最後にlist()を用いてリストに変換してやるようにする。

# D': スペースで区切られた2つの整数N, Mを受け取る
# 3 4
N, M = map(int, input().spilt())

対して、複数の変数に一度に代入を行っている（アンパック代入と呼ばれる）例D’では、右辺のmap objectは自動的に展開（アンパック）されるため、listを使用する必要はない。もちろんリストに変換しても動作するが、わずかに無駄な書き方になる。

# E: 最初に要素数Nが指定され、続くN行の整数を受け取る
# 3
# 1
# 2
# 3
N = int(input())
integer_list = [int(input()) for _ in range(N)]

例Eでは、リスト内包表記を使用している。この部分を、forを用いて分かりやすい形に書き下すと、次のようなものになる。

integer_list = []
for i in range(M):
    n = int(input())
    integer_list.append(n)

リスト内包表記は非常に強力で、また、慣れると書くのも早いので、一般的な使い方も覚えておくと便利だ。

そして、これらを組み合わせて入力を扱ったのが例Fということになる。このパターンでは、2次元リストとして入力を受け取ることになる。

# F: 最初に要素数N、行数Mが指定され、続くM行の整数列（要素数N）を受け取る
# 2 3
# 1 2
# 3 4
# 5 6
N, M = map(int, input().split())
integer_matrix = [list(map(int, input().split())) for _ in range(M)]

以上のパターンが理解できれば、これらを組み合わせることで、基本的な入力については問題なく取り扱えるだろう。