AtCoderのまあまあ便利なPython標準入力フォーマット

まあまあ便利です。

基本

read.py

import sys

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline

こいつらで入力すると速い。inputよりもこちらを使う方が実行時間で有利です。

readとreadlineの違い

readはEOFまで入力を受け取るのに対して、readlineは各行ごとに入力を受け取ります。AtCoderのコードテストなどではEOFは気にしないでもいいですが、手元実行の場合は入力後Ctrl+Dを押して、入力が終わったこと（EOF）を教えてあげる必要があります。また、EOFの性質上readは一番最後にしか使えません。よって、以下のようなことができません。

a = read()      #どこまで見ればいいのか分からない！
b = readline()  

入力の加工

また、このコードだと受け取った入力はバイト型で、しかも改行コードを含んでいます。

a = readline()  # abc　を入力
b = read()  　　　　　　　# a b
            　　　　　　　# c d を入力

print(a)  # 結果 b'abc\n'
print(b)  # 結果 b'a b\nc d\n'

文字列として受け取りたいときはこのままだと扱いにくいので、decodeしてrstripで改行コードを消してあげます。また、readならばdecodeしてsplitすると、文字列リストを作ることができます。これはsplitが標準では' 'や'\n'といった空白文字ごとに文字列を区切ってリストに変換するからです。

a = readline().decode().rstrip()  # abcを入力
b = read().decode().split()       # a b
                                  # c d を入力

print(a) # 結果 abc
print(b) # 結果　['a', 'b', 'c', 'd']

標準入力用関数の作成

上記のコードの下に、listを使うかNumPyのndarrayを使うかに応じて、以下のいずれかを書きます。

read1.py

#listで受け取る場合

def from_read(dtype=int):
    return [dtype(x) for x in read().decode().split()]

def from_readline(dtype=int):
    return [dtype(x) for x in readline().decode().split()]

def read_matrix(N, M, dtype=int):
    matrix = [None for i in range(N)]
    for i in range(N):
        matrix[i] = from_readline(dtype)
    return matrix

read2.py

#ndarrayで受け取る場合

def from_read(dtype=np.int64):
    return np.fromstring(read().decode(), dtype=dtype, sep=' ')

def from_readline(dtype=np.int64):
    return np.fromstring(readline().decode(), dtype=dtype, sep=' ')

def read_matrix(N, M, dtype=np.int64):
    matrix = np.empty((N, M), dtype=dtype)
    for i in range(N):
        matrix[i] = from_readline(dtype)
    return matrix

dtypeはデータの型でint, strなどを意図していますが、本来はNumPyの用語なので注意。list, ndarrayで使い心地が変わらないようにするために引数名を揃えています。（list版read_matrixに使いもしないMを入れているのはそのため）
ただし、ndarray版ではstrは使えません。あくまで数値計算のために使いましょう。

各関数の説明

• from_read - EOFまでlist(ndarray)に要素を格納
• from_readline - 改行までlist(ndarray)に要素を格納
• read_matrix - (N, M)の行列をlist(ndarray)として受け取る

　list（しかりndarray）の入力関数を用意しておくことのメリットはまあまあ大きく、AtCoderにおいてはほとんどの問題で役立ちます。

例えば、数列の受け取りにそのまま使えるのはもちろん、unpackによって要素を個別に代入できます。

実際の受け取り例

import sys

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline

def from_read(dtype=int):
    return [dtype(x) for x in read().decode().split()]

def from_readline(dtype=int):
    return [dtype(x) for x in readline().decode().split()]

def read_matrix(N, M, dtype=int):
    matrix = [None for i in range(N)]
    for i in range(N):
        matrix[i] = from_readline(dtype)
    return matrix

X = int(readline())     # 100
N, M = from_readline()  # 2 3
A = from_readline()     # 3 1 4 1 5
B = from_read(str)      # a b
                        # c d e

print(X)  # 100
print(N)  # 2
print(M)  # 3
print(A)  # [3, 1, 4, 1, 5]
print(B)  # ['a', 'b', 'c', 'd', 'e']

使いこなすとまあまあ快適です。