0. はじめに
「Python3を少し知っている。競技プログラミングを始めたい。Python3で競技プログラミングを始めよう」という方向けの記事です。map(), iter(), open(0) などを使った、知っておくと便利な書き方も紹介します。
Python2の方はPythonで競技プログラミングする時に知っておきたいtips(入出力編) などが参考になると思います。
出題例のリンクを貼り付けた箇所があります。問題は、AtCoder に登録したら次にやること ～ これだけ解けば十分闘える！過去問精選 10 問 ～ を参考に選びました。
1. まとめ
競技プログラミングで使える Python3 の入力方式は3種類1あります。
input()sys.stdinopen(0)
input() だけで、次に示す関数やメソッドなどを組み合わせることであらゆる入力形式に対応できます。
大量に1行ごとの入力を行う場合は sys.stdin.readline() を使うと、多少速くなることがあります2。また、読み取る変数の数が不明で簡単に書けます。 ただし、import sys の書き忘れに注意です。改行に左右されない書き方をしたいときは open(0) を使うことで実現できます。
上記のどれを使っても、はじめ文字列として読み込まれます。整数や小数として読み込みたい場合は、変換する必要があります。
- 整数に変換したいとき ->
int() - 小数に変換したいとき ->
float()
さらに問題の形式に応じて、以下の3つを駆使しましょう。
- スペース区切りのとき ->
.split()（+map()） - 複数行を取り込みたいとき -> 「リスト内包表記」
- 特定の文字列が来るまで読み込み続けたいとき ->
iter()
以下、問題の形式別に説明していきたいと思います。細かく分けていたら長くなってしまいました。パターン別に覚える必要はありません。
2. input()
入力から1行を文字列として読み込みます。末尾の改行は除かれます。
2-1. 1行で1つの入力
2-1-1. 「1行に文字列1つ」
出題例: 「過去問精選10問」の第9問 ABC 049 C - Daydream
形式
dreamerer
コード
S = input() # type(S) => <class 'str'>
2-1-2. 「1行に整数1つ」
int() を使います。
出題例: ABC 104 A - Rated for Me
形式
1200
コード
N = int(input()) # type(N) => <class 'int'>
2-1-3. 「1行に小数1つ」
float() を使います。
形式
3.14
コード
x = float(input()) # type(x) => <class 'float'>
2-2. 1行に複数の入力
主に .split() を使います。
str.split(sep) は、str を sep を区切り文字列として分割した単語のリストを返します。sep を指定しなければ、空白を区切り文字列として分割してくれます。
str.split(sep)の例
>>> '1,2,3'.split(',')
['1', '2', '3']
2-2-1. 「1行に複数の文字列」
出題例: ABC 059 A - Three-letter acronym
形式: 空白区切りで３つの文字列の場合
atcoder beginner contest
コード例1: リストとして受け取る
S = input().split() # type(S) => <class 'list'>
print(S) # ['atcoder', 'beginner', 'contest']
コード例2: 文字列として受け取る（変数の数が既知のときのみ）3
A, B, C = input().split() # type(A) => <class 'str'>
print(A) # 「atcoder」と表示される
2-2-2. 「1行に決まった数の文字」
文字数がわかっている文字列を、1文字ずつ変数に格納したいときに使えます。
出題例: ABC 095 A - Something on It
形式: 3文字の場合
oxo
コード
a, b, c = input() # 3文字の文字列
print(a) # 「o」と表示される
参考: 文字列として変数に格納し、インデックス指定を行う
S = input()
print(S[0]) # 「o」と表示される
2-2-3. （応用） 文字列を集合として読み込みたい
set() を使います。
形式
atcoderregularcontest
コード:
S = set(input()) # type(S) => <class 'set'>
S = {*input()} と書くこともできますが、AtCoder（Python 3.4.3）では SyntaxError となります。
2-2-4. 「1行に複数の整数」
map() と list() を使うか、「リスト内包表記」で書きます。
出題例: 「過去問精選10問」の第3問 ABC 081 B - Shift only
第6問 ABC 088 B - Card Game for Two
形式
$N$
$A_1\ A_2\ \dots\ A_N$
4
5 6 8 10
コード例1: map() と list() を使う
map()の使用には、多少注意が必要です。map() を使う際に気をつけたいことを末尾に付けました。
N = int(input())
A = list(map(int, input().split())) # 入力のときに N を使う必要はありません
print(A[0]) # 「5」と表示されます
A = [*map(int, input().split())] と書くこともできますが、AtCoder（Python 3.4.3）では SyntaxError となります。
コード例2: 「リスト内包表記」で書く
N = int(input())
A = [int(x) for x in input().split()]
2-2-5. 「1行に決まった数の整数」
出題例: 「過去問精選10問」の第1問 ABC 086 A - Product
第5問 ABC 083 B - Some Sums
第8問 ABC 085 C - Otoshidama
形式: 変数が２つの場合
3 4
コード
a, b = map(int, input().split()) # list() を使っても構いません
参考: map() を使わずに書くこともできます
ab = input().split()
a, b = int(ab[0]), int(ab[1])
2-2-6. （応用） 「1行に複数の整数」をソートしておきたい
形式
4
20 18 2 18
コード例1: map() を使う
list() を使う必要はありません。代わりに sorted() を使います。sorted() は、[2, 18, 18, 20] のように昇順にソートされたリストを返す関数です。
N = int(input())
A = sorted(map(int, input().split())) # type(A) => <class 'list'>
コード例2: 「リスト内包表記」で書く
N = int(input())
A = sorted([int(x) for x in input().split()])
A = sorted(int(x) for x in input().split()) のように [ ] を省略しても構いません。
2-2-7. （応用） 「1行に複数の整数」を集合として読み込みたい
形式
4
20 18 2 18
コード例1: map() を使う
list() を使う必要はありません。代わりに set() を使います。
N = int(input())
A = set(map(int, input().split())) # type(A) => <class 'set'>
A = {*map(int, input().split())} と書くこともできますが、AtCoder（Python 3.4.3）では SyntaxError となります。
コード例2: 「セット内包表記」で書く
N = int(input())
A = {int(x) for x in input().split()}
# A = set([int(x) for x in input().split()]) と書くこともできる
2-3. 複数行の入力
各行の入力方法（「2-1 1行で1つの入力」と「2-2 1行に複数の入力」）。
2-3-1. 最初にすべて読み込む必要がある場合 - 行列、グリッド
出題例: ABC 075 B - Minesweeper
「過去問精選10問」の第7問 ABC 085 B - Kagami Mochi
形式
$H$ $W$
$S_1$
$\vdots$
$S_H$
コード例1: 配列を用意して、追加していく
H, W = map(int, input().split())
S = []
for _ in range(H):
S.append(input()) # S += [input()] とも書ける
整数として取り込むときは、 input() を int(input()) に置き換えます。
コード例2: リスト内包表記
H, W = map(int, input().split())
S = [input() for _ in range(H)]
# S = [input() for _ in [0] * H] とも書ける
2-3-2. 最初にすべて読み込む必要がある場合 - 事前に各リストを用意しておく
出題例: 「過去問精選10問」の第10問 ABC 086 C - Traveling
形式
$N$
$t_1\ x_1\ y_1$
$t_2\ x_2\ y_2$
$\vdots$
$t_N\ x_N\ y_N$
2
3 1 2
6 1 1
コード
N = int(input())
t = [0] * N
x = [0] * N
y = [0] * N
for i in range(N):
t[i], x[i], y[i] = map(int, input().split())
2-3-3. 1行ずつ処理 - for 文を使うだけ
出題例: ABC 061 B - Counting Roads
形式
$N\ M$
$a_1\ b_1$
$\vdots$
$a_M\ b_M$
4 3
1 2
2 3
1 4
コード
N, M = map(int, input().split())
for _ in range(M):
a, b = map(int, input().split())
# 処理
2-3-4. 特定の文字列を読み込むまで読み込み続ける
while と if を用いて書くこともできますが、ここでは iter を使ってみたいと思います。
AOJで見かける形式です。
形式: "-1"が入力の終わりを示す場合
$a_1$
$a_2$
$\vdots$
$-1$
コード例
for e in iter(input, '-1'): # 「'-1'」の部分を問題に応じて変更します
# 処理
e は文字列なので、整数として扱いたいときは int() を使いましょう。
3. sys.stdin
3-1. 入力が終わるまで読み込み続ける
input() を使う場合は try-except を使います。ここでは、input() を使うより簡単に書けるコードを紹介します。
AOJで見かける形式です。
形式
$a_1$
$a_2$
$\vdots$
(EOF)
コード
import sys
for e in sys.stdin:
# 処理
e は文字列なので、整数として扱いたいときは int() を使いましょう。また、e の末尾には改行コードが含まれているので注意しましょう。
各行が整数の場合は、in map(int, sys.stdin) と書くこともできます。
3-2. input() の代わりとして使う - .readline()
まず、sys モジュールを読み込みます。
import sys
3-2-1. １つ文字列を読み込む
S = sys.stdin.readline().rstrip()
末尾の改行コードを取り除くために、rstrip メソッドを使いました。
S = sys.stdin.readline()[:-1]
最後の一文字を取り除ければよいので、どんな方法でも構いません。
3-2-2. その他の形式1つの整数を読み込む
今までに紹介したコードの input() を sys.stdin.readline() に書き換えるだけです。
コード例1: 1つの整数
N = int(sys.stdin.readline())
コード例2: 複数の文字列
X = sys.stdin.readline().split()
コード例3: 複数の小数
A = [float(x) for x in sys.stdin.readline().split()]
# A = list(map(float, sys.stdin.readline().split())) とも書けます
3-3. 入力のリストとして扱う - .readlines()
出題例: ABC 106 D - AtCoder Express 2
形式
$N$ $M$ $Q$
$L_1$ $R_1$
$L_2$ $R_2$
$\vdots$
$L_M$ $R_M$
$p_1$ $q_1$
$p_2$ $q_2$
$\vdots$
$p_M$ $q_M$
コード
import sys
s = sys.stdin.readlines()
N, M, Q = map(int, s[0].split())
for e in s[1:M + 1]:
L, R = map(int, e.split())
# L, R に関する処理
for e in s[M + 1:]:
p, q = map(int, e.split())
# p, q に関する処理
4. open(0)
4-1. 決まった数の整数
出題例:「過去問精選10問」の第4問 ABC 087 B - Coins
形式1
$A$ $B$ $C$ $X$
形式2
$A$
$B$
$C$
$X$
4-1-1. input() を使う
input() を使うと、たとえば、
形式1 の場合
A, B, C, X = map(int, input().split())
形式2 の場合
A, B, C, X = [int(input()) for _ in range(4)]
のように、コードを書き分ける必要があります。
4-1-2. open(0) を使う
しかし、open(0) を使えば、
A, B, C, X = map(int, open(0).read().split())
と同じコードで、形式1 にも 形式2 にも対応することができます。
4-2. 複数の整数
変数の前に * をつけます4。
出題例:「過去問精選10問」の第7問 ABC 085 B - Kagami Mochi
形式
$N$
$d_1$
$\vdots$
$d_N$
4
10
8
8
6
コード
N, *d = map(int, open(0).read().split()) # type(d) => <class 'list'>
リストの場合は * をつけます。1つ目の整数が N に代入され、2つ目以降はすべて、 d にリストとして代入されます。
ただし、Nも含めて1つの変数に代入する場合はカンマ忘れに注意です。
*X, = map(int, open(0).read().split()) # カンマ(,)が必要
print(X[0]) # => 4
print(X[1:]) # => [10, 8, 8, 6]
5. 補足
5-1. map() について
map() は、リストのすべての値に値に対して関数を適用する組み込み関数です。
map(function, iterable) のように使います。map オブジェクトを返します。
map()の例
>>> nums = [1, 4, 9, 16, 25]
>>> sqrt = lambda x: x ** 0.5
>>> list(map(sqrt, nums))
[1.0, 2.0, 3.0, 4.0, 5.0]
5-2. map オブジェクトを扱うときに注意すること
map オブジェクトはリストと扱いが異なる場合があるので、注意が必要です。
list() を使って、リストに変換することで解決します。
リストと同様に扱えるケースとして、for 文があります5。
>>> A = ['1', '2', '3']
>>> for x in map(int, A):
... print(x ** 3)
...
1
8
27
リストと扱いが異なる点として例えば、インデックス指定ができません6。
>>> A = map(int, input().split())
5 6 8 10
>>> A[0]
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: 'map' object is not subscriptable
また、一度iterateしてしまうと、再利用ができません。
>>> A = map(int, '5 6 8 10'.split())
>>> B = list(A)
>>> print(B)
[5, 6, 8, 10]
>>> C = list(A)
>>> print(C)
[]
-1. おわりに
加筆して、競技プログラミングの入力形式の網羅に挑戦してみたいと思います。
map や * といった Python3 の文法は、別の記事として書き直すかもです。