Pythonで競技プログラミングする時に知っておきたいtips(入出力編)

  • 283
    いいね
  • 11
    コメント

Pythonで競技プログラミングする時に知っておきたいtipsの,入出力の部分を分割しました.

Pythonで競技プログラミングをする際の入出力手法について例示する.
入出力のコードは,言語を初修する際のハードルになることが多いが,一度覚えてしまえばワンパターンなのでさっさと覚えてしまいたい.

参考

標準入力の受け取りのその1 - Qiita

1行に文字列1つの入力

形式

input
S

Sは文字列.

入力

S = raw_input()

1行に整数または浮動小数点数1つの入力

形式

input
N

Nは整数または浮動小数点数.

入力

N = input()

1行にスペース区切りで複数の文字列の入力

形式

input
A B C

A,B,Cは,いずれも文字列.

入力

l = raw_input().split()
print l[0] # Aの値が出力される
print l[1] # Bの値が出力される
print l[2] # Cの値が出力される

split()は,スペース区切りで文字列をリストに分割する.

1行にスペース区切りで複数の整数または浮動小数点数の入力

形式

input
A B C

A,B,Cは,いずれも整数または浮動小数点数.

入力

すべての変数をリストに代入したいとき

A,B,Cがすべて整数の場合は,

l = map(int, raw_input().split())
print l[0] # A
print l[1] # B
print l[2] # C

要するに,文字列型で受け取った入力をスペースで分割してリストに代入し,各要素を整数型にmapするということ.

要素の数が入力データに与えられることがあるが,Pythonの場合は使わなくてもさほど問題にならない.

浮動小数点数の場合は,

l = map(float, raw_input().split())

浮動小数点数の場合にintを使うと,値が切り捨てられてしまうので注意.

それぞれを別の変数に格納したい場合

A, B, C = map(int, raw_input().split())
print A # A
print B # B
print C # C

リストの要素数が左辺の数に一致していれば,左辺それぞれの変数にリストの各要素を代入できる.
こういう形式の代入を「アンパック代入」というらしい.

競技プログラミングの場合,入力の要素数が明らかで,かつ少ない場合のみ使うとよい.

あまり見かけないが,1行に文字列や整数が混在している場合は,とりあえず1行を文字列として受け取り,後でint()float()などを用いて変換する.

複数(N)行それぞれに1つの整数の入力

形式

input
N
a1
a2
a3
...
aN

a1,a2,...,aNは,N個の整数.

入力

N = input()
a = []
for i in range(N):
    a.append(input())
print a # [a1, a2, a3, ..., aN]

これでいい.

@zakuro9715さんに教えていただきました.

このコードは,リスト内包表記を用いて,以下のようにも表せる.
N行入力が1行で書けてキレイ.

N = input()
a = [input() for i in range(N)]
print a # [a1, a2, a3, ..., aN]

複数行(特定の整数を受け付けるまで)それぞれに1つの整数の入力

Aizu Online Judge (AOJ)とかでよく見る形式.

形式

input
a1
a2
a3
...
-1

a1,a2,...,はいずれも整数.
-1が入力の終わりを表す.

入力

a = []
while True:
    n = input()
    if n == -1:
        break
    a.append(n)
print a # [a1, a2, a3, ..., aN]

複数行(入力の終わりまで)それぞれに1つの整数の入力

これもAOJでたまに見る.
デバッグもしづらいし,個人的には嫌いな形式.

形式

input
a1
a2
a3
...
(EOF)

a1,a2,...,はいずれも整数.
EOFが来るまで読む.

入力

import sys

a = []
for line in sys.stdin:
    a.append(int(line))
print a # [a1, a2, a3, ...]

複数(N)行それぞれにM個の整数の入力

今までのまとめ的な感じ.

形式

input
N M
a11 a12 a13 ... a1M
a21 a22 a23 ... a2M
a31 a32 a33 ... a3M
...
aN1 aN2 aN3 ... aNM

a11,a12,...,aNMは整数.

入力

N, M = map(int, raw_input().split())
a = []
for i in range(M):
    a.append(map(int, raw_input().split()))
print a 
# [[a11, a12, a13, ..., a1M]
#  [a21, a22, a23, ..., a2M]
#  ...
#  [aN1, aN2, aN3, ..., aNM]]

たぶん,これぐらいおさえておけば困らない.

ジェネレータを使った汎用的な入力

@t2y さんに教えていただきました.

ジェネレータを用いることによって,複数行に渡る入力をうまく受け取ることができる.

例えば,

形式

input
a1
a2
a3
...
(EOF)

a1,a2,...,はいずれも文字列.
EOFが来るまで読む.

のような形式であれば,

def get_input():
    while True:
        try:
            yield ''.join(raw_input())
        except EOFError:
            break

if __name__ == '__main__':
    a = list(get_input()) # [a1, a2, a3, ...]

のようにする.

他にも,get_input()の内容を,

def get_input():
    while True:
        try:
            s = raw_input()
            if s == '-1':
                break
            yield ''.join(s)
        except EOFError:
            break

と変えることで,「"-1"を読むか入力の終わりまで」といった若干複雑な入力を受け取るといったこともできる.

出力

1変数を出力(改行あり)

形式

output
A

末尾に改行を含む.

出力

print A

整数型でも文字列型でも出力できる.
他にも,リストやディクショナリなどのデバッグもprint文でそのままできる.

1変数を出力(改行なし)

正直見たことがないが,一応勉強のために.

形式

output
A

末尾に改行を含まない.

出力

import sys

sys.stdout.write(A)

複数変数をスペース区切りで出力

これはよく見る.

形式

output
A B C

末尾に改行を含む.

出力

print A, B, C

print文では,変数の後に,を入れると,改行が半角スペースに変換される.

Python Tips:改行なしで文字列を出力したい - Life with Python

複数変数をカンマ区切りで出力

形式

output
A,B,C

A,B,Cはいずれも文字列.
末尾に改行を含む.

出力

print ','.join([A, B, C])

join()は,split()の逆の演算に相当し,引数の文字列のリストの各要素をレシーバの文字列で結合したものを返す.

リストのすべての要素が文字列型でなければならないことに注意.

ここで,例えば,A,B,Cが整数だったとすると,

print ','.join(map(str, [A, B, C]))

のように,str()によって各要素を文字列に変換してからjoinする.

フォーマットにしたがった出力

C++におけるprintf関数のように,自由な形式での出力がしたい場合がある.

形式(例)

output
Case 1: A B

A,Bはいずれも整数.
末尾に改行を含む.

出力

print 'Case 1: {0} {1}'.format(A, B)

format()は可変長の引数を受け取り,i番目の引数の値を'{i}'に格納した文字列を返す.

出力時の基数変換や桁揃えなど,printfでできるだいたいのことがformat()でもできる.
詳しくは,リファレンスを見るとよい.