Python 競プロ用 Tips

既に巷に溢れている気もするが、まぁリンク集+自分用メモ的な感じで。

Pythonで競プロに参加する上での、言語依存な書き方など。散発的になるのは仕方なし。

AtCoderにおけるPythonの環境

2023年12月アップデート後

古い情報

• Python3で競技プログラミングする時に知っておきたいtips（入力編） - Qiita

主に4つの方式がある

• input()
• sys.stdin 系
• sys.stdin.buffer 系
• open(0)
  • あまり使ったことがない。コードゴルフとかでたまに見る。sys.stdin系とほぼ一緒

基本。1行読み込んだ文字列が得られる。末尾の改行は取り除かれる。

s = input()
print(s)
print(type(s))

# 入力
# 3 1 4 1 5

# => '3 1 4 1 5' <class 'str'> 型

入力が大量で、速さを求めるとき。

• sys.stdin
  • 1行ずつ読み込むイテレータが得られる。for line in sys.stdin: 等と使う
• sys.stdin.readline()
  • 1行読み込んだ文字列が得られる
• sys.stdin.readlines()
  • 残っている入力を全て読み、1行1要素の文字列のリストにする
• sys.stdin.read()
  • 残っている入力を全て1つの文字列として読み込む

注意点

• input() と異なり、末尾の改行は残る
  • 数値になおす場合、int(), float() などは末尾に改行が付いていても解釈できるので、特に取り除く必要は無い
  • そのまま文字列として扱う場合に、改行文字に注意
  • s.rstrip() で、sについた末尾の改行は削除できる
• 入力の終了
  • 入力の終わりは「EOF」という特殊な記号で判断するので、手元で入力を直接打ち込むなどで与える場合、そのままでは実行が始まらない
  • Ctrl-D で「EOF」を入力できる

import sys

s = sys.stdin.readline()
print(s)
print(type(s))

t = sys.stdin.readlines()
print(t)
print(type(t))

# 入力
# 3 1 4 1 5
# 9 2 6 5 3
# 5 8 9 7 9

# s: '3 1 4 1 5\n'                   <class 'str'>
# t: ['9 2 6 5 3\n', '5 8 9 7 9\n']  <class 'list'>

なるべく read() で一括で読み、一括で split() や map(int, ○○) に渡した方が速い。

N M
A1 A2 A3 ... AN
B1 B2 B3 ... BM

のようなデータなら、アスタリスク記法を使って以下のようにすると、簡潔にA,Bを整数リストで得られる。

• for-unpacking-the-containers-to-variables.py

n, m, *ab = map(int, sys.stdin.read())
aaa = ab[:n]
bbb = ab[n:]

# 入力
# 5 4
# 1 2 3 4 5
# 6 7 8 9

# ab:  [1, 2, 3, 4, 5, 6, 7, 8, 9]
# aaa: [1, 2, 3, 4, 5]
# bbb: [6, 7, 8, 9]

少し困るのが、「$N$ 行に渡って $a_i,b_i,c_i$ 3種類の整数が与えられる」みたいな時、一括で読んでしまうと1行ずつの情報としてまとめにくい。

N
a1 b1 c1
a2 b2 c2
...
aN bN cN

以下のようにすると、簡潔に [(a1,b1,c1), (a2,b2,c2), …] にできる（ちょっとテクニカル）。1行ずつ読んでリストに追加していくより1.5倍程度速い。

n = int(sys.stdin.readline())
mp = map(int, sys.stdin.read().split())
abc = list(zip(mp, mp, mp))  # ←1行に含まれるデータ数だけmpを連ねる

# 入力
# 3
# 1 2 3
# 4 5 6
# 7 8 9

# abc: [(1, 2, 3), (4, 5, 6), (7, 8, 9)]

実験コード

import sys
from time import time


# 紹介した方法、mapをzip化
def zip_map():
    start = time()
    n = int(sys.stdin.readline())
    mp = map(int, sys.stdin.read().split())
    ab = list(zip(mp, mp))
    end = time()
    print(end - start)


# 一旦listにして、1つ飛ばしでsliceしたのを再度zip化
def slice():
    start = time()
    n = int(sys.stdin.readline())
    ab = list(map(int, sys.stdin.read().split()))
    ab = list(zip(ab[0::2], ab[1::2]))
    end = time()
    print(end - start)


# 1行ずつ読んでlistにappend
def eachline():
    start = time()
    n = int(sys.stdin.readline())
    ab = []
    for line in sys.stdin:
        ab.append(tuple(map(int, line.split())))
    end = time()
    print(end - start)

# N=200000 の入力を与えてテスト

# zip_map()  # 0.17016100883483887
# slice()  # 0.19611454010009766
# eachline()  # 0.26198315620422363

より大量の入力を高速にさばきたい場合。ただし癖が強いので十分な理解が必要。

• maspyさんのツイート: "pythonで競プロをやる情報をあさると、入力の受け取り方で sys.stdin.readline (read, readlines) が挙げられていることが多いが、sys.stdin.buffer.readline (read,readlines) の方が多少速いことが多い。特にpypyだと差がつきがち。"

それぞれの役割は、sys.stdin.○○ と同じ。末尾に改行が残るのも同じ。

ただし、以下に示す注意点がある。

buffer.read系を用いる場合は、事前に input や sys.stdin.read系を使っていてはいけない。

input等を使って読み込んだ時点で、既にある程度先までバッファで読まれているため、あらためて buffer.read系で読んだとき、中途半端な位置からの読み込みとなってしまう。

buffer.read系を用いる場合は、最初からそれのみで完結させる。

まず、Pythonの文字列（っぽいもの）を表す型として、str型の他に、bytes型（バイナリシーケンス型）がある。

str型はPython3内部での文字を扱う形式であるUnicode型に変換された上での文字列情報だが、bytes型は、文字コード不定の、まさに入力バイナリ列そのものである。

• 組み込み型 — Python 3.8.3rc1 ドキュメント
• Python3で文字列を処理する際の心掛け - Qiita

このbytes型は、

• 文字コードがasciiであることを前提とした上で、だいたいstr型と同じように扱える
  • s.split() はできる
  • int(), float() に与えると適切に数値化してくれる
• 文字列比較は型を揃える必要がある（前もってdecodeしてstr型にするか、対象もbytes型にするか）
• print(s) は、一応できる
  • 「b'abcd'」というようにbytes型であることを示す b' ' が付く
  • ascii外の文字が含まれていた場合、文字化けする

なので、実質困ることがあるというと、入力に文字列が含まれるときである。 sys.stdin でも十分高速なので、文字列が含まれる場合は無理せずそちらを使った方が面倒がない。

どうしてもsys.stdin.bufferを使いたい場合、（競プロでは普通、英数字しかないので）decode('ascii') 等を指定してデコードすれば、str型となる。

• collections --- コンテナデータ型 — Python 3.8.3rc1 ドキュメント

• 紹介
  • scipyのFloyd-WarshallとDijkstraのすすめ　Python 競技プログラミング Atcoder - じゅっぴーダイアリー
• グラフの作り方
  • SciPyによる経路探索

実際は、探索しながら他の値も計算しなければならない、とかで、経路探索であってもこれをそのまま使えない問題も多い。

典型的な問題では使えばよいが、一応自分でアルゴリズムを書けることも重要。

• scipy.spatial.ConvexHull — SciPy v1.4.1 Reference Guide

たまに出る幾何で、与えられた2次元座標の点を全て包む最小の多角形を求める問題。

経路探索と同様、なかなかそのまま使える機会は少ないが、一応できるということを知っておくだけでも。

多くの競プロサイトはNumbaは使えないことが多いが、AtCoderは2020/04より使えるようになった。

そもそも、Pythonはお手軽に書けるのが一つの利点であり、 使える関数やデータ型が限定されるNumbaに書き換えてまで Pythonにこだわる意味があるのか、というのはある。
AtCoder以外では今のところ使えないし、複雑な処理ではコンパイル可能なコードにするのが難しいこともある。

そうは言っても、メイン言語がPythonの人にとって、1からC++などで書きなおす以外にPythonで通せる幅が広がったのは歓迎すべきことだ。

Numbaの中でも使い方は複数あるが、競プロの文脈でポイントとなるのは、以下だろう。

• 「実行時間」を短くしたい
  • 高速化の恩恵の大きい “NoPython” モードでのコンパイルが基本となる
  • 実行時コンパイル（JIT）では、実行時間に解析・コンパイル時間が含まれるため、なるべく事前コンパイル（AOT）したい
• 使わない方がよい問題もある
  • 文字列の方が扱いやすい問題とか、多倍長整数が有効な問題とか
  • →全ての問題に共通して使うテンプレート、というのではなく、必要に応じて呼び出せるスニペットがよい
• ローカル環境でテストした上で提出する
  • ローカル用と提出用でコードの変更が必要なら、ローカル用のコードをうっかり提出してしまうケアレスミスが発生する
  • →なるべく同一のコードで動く方がよい

このことを念頭に置いて、スニペットを作成する。

以下、考慮すべきことと採択理由をメモしている。

Numbaのコンパイルには、大別して実行時コンパイル（JIT）と事前コンパイル（AOT）がある。
JITはその名の通り実行時間にコンパイル時間が含まれてしまうので、基本的にはAOTを行いたい。

ただし、JITでも cache=True（キャッシュを有効にする）オプションを指定でき、 それを用いると自動的にコンパイル結果をファイルとして保存し、以前のコンパイル結果があればそれを利用することができる。実質AOT。 ¹⁾

• Notes on Caching — Numba 0.50.1 documentation

つまり、以下のいずれかを行えばよい。

• AOTコンパイル
• cache=True としたJITコンパイル

一長一短あるので使い分け。でも基本的にはJITの方が使い勝手がよい。

• コードの見やすさ・改修の容易さ
  • AOTは、コンパイルと本処理で処理が明確に異なるため、切り分ける必要があり、記述が多少ごちゃごちゃしてしまう
  • JITは関数に @njit デコレータを付けるだけでコンパイル・本処理で共通のコードで動く
• ローカル環境での実行
  • ローカル環境でAOT用のコードをジャッジシステムと同様に動かそうと思えば、コンパイルと本処理で2回実行することになる。それは面倒なので他の方法を採るにしろ、やはり多少、専用の記述が必要になる
  • JITは、ローカルでも提出コードで同じように動く
• 実行速度
  • AOTは、本処理時にはコンパイルしたモジュールのみ読み込めばよい
  • JITはNumbaライブラリ読み込みに加え、キャッシュの存在チェックなどが必要となり、環境にも依るが100ms程度は余分にかかる

コードをスッキリと保ちたいならJIT、ギリギリまで速度を求めたければAOT、でよいと思う。

こう書くとAOTはやたら面倒そうだが、スニペットとしてまとめ、使い方さえ把握しておけば大丈夫。

AOTの方を採用する場合でも、なるべく提出コードを変えずにローカルで動くようにしたい。

AOTを用いる場合、ジャッジシステムでの挙動と同様にコンパイルと本処理を切り分けると、2回実行する必要があり、煩わしい。 また、1つのコンパイル結果をそんなに何回も使い回さない。

それなら、毎回数秒かかってしまうが、常にコンパイル→本処理と1回で流すようにする方が楽。
さらに、それならローカルではAOTでなくJITでいいよね、という話になる。

実際、AOTをローカル環境でやろうとすると、C++コンパイラが必要、JITよりも若干コンパイル時間がかかる、などあるので、手元ではJITの方で十分と感じる。

OS環境はWindowsなら os.name == 'nt'、Mac/Linuxなら os.name == 'posix' で、ローカルがWindowsであればジャッジシステムと区別できる。 これで切り分けて、ローカルならJITコンパイルして実行するようにスニペットを作っておく。

Numba関数内ではinput()など入力受け取り関数は使えない。素のPythonで受け取って、引数として与える必要がある。

そして、関数をコンパイルするには、引数と戻り値の型を指定してやる必要がある。

• Types and signatures — Numba 0.50.0.dev0+236.g64fbf2b-py3.7-linux-x86_64.egg documentation

厳密には、JITの場合は型指定せずとも動くのは動く。 ただしその場合、実際に関数が呼ばれたタイミングで自動推論→コンパイルされるのであり、逆に言えば呼ばれるまでされない（本来の意味での実行時コンパイルになる）。 型指定すれば定義された時点でコンパイルされるので、コンパイルフェーズに行いたければやはり型を指定する必要がある。

毎回きちんと型指定するのは手間だし、どう書けばいいんだっけ、と間違いも起こりやすい。

AtCoderでは、文字列を含まない問題の入力は全て整数であることが多い。 そのため、以下のようにすると、一括で全入力を1つのNumpy配列にできる。

inp = np.fromstring(sys.stdin.readline(), dtype=np.int64, sep=' ')

これを唯一の引数とするようにNumba関数を作ると、型指定は常に '(i8[:],)' で済む。戻り値は省略しても自動推論される。

問題に合わせて変更する箇所が少なくなり、「Numba関数の中身だけ書けばいい」という状態にできる。

小数や文字列を含む場合などは、共通化しにくいのでその都度、型を適切に合わせることにする。

以下を参照させていただいた。あまり変わっていないが、ローカル環境でも同一コードで動くようにした。

• AtCoderで Python を高速化する Numpy + Numba を使う - Qiita

ローカル環境がWindowsであることを前提としており、os.name でLinux（AtCoderのジャッジシステム）と区別できる。 MacやLinuxなら、platform モジュールなどでジャッジシステムと違う箇所を探すか、ローカル環境のみ実行時オプションを付けて識別するなどする。

基本的に solve() の中身と、あと必要なら最後の出力整形部分のみ書けばよい。

実装例

import os
import sys

import numpy as np


def solve(inp):
    a, b, c = inp
    return a + b + c


def cc_compile():
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', '(i8[:],)')(solve)
    cc.compile()


if sys.argv[-1] == 'ONLINE_JUDGE':
    cc_compile()
    exit()

if os.name != 'posix':
    from numba import njit
    
    solve = njit('(i8[],)', cache=True)(solve)
else:
    # noinspection PyUnresolvedReferences
    from my_module import solve

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)

処理の切り出しなどで関数を複数に分けたくなったとき、以下の2つがあるが、

• それぞれをグローバルに書いて個別にコンパイル
• 1つの大枠の関数の中で個々の関数定義を書いて、大枠の関数のみコンパイル（内部関数）

前者は関数毎にコンパイル指定（JITなら@njitデコレータの付与、AOTならcc.exportへの登録）の必要がある一方、 後者は大枠関数のみの指定で特に問題なく内部の関数も型推論してくれる。

ただし、内部関数が再帰を含む場合、コンパイルが通らない。

numba.core.errors.NotDefinedError: Variable '(内部関数名)' is not defined.

通常は1つの大枠関数に入れた方が手間が少ないのでそうし、再帰関数のみ個別にコンパイルする。

@jitclass デコレータによりJITでのコンパイルは可能だが、これはこれで関数とは別の書き方が必要であり、覚えることとが増えてしまう。

AOTでのコンパイル方法は探したけど見つかってない。

クラスは、1つのオブジェクトに関係する処理をまとめることで理解しやすくしてくれる点はあるが、競プロのような短いコードでは必須でもないので、今のところはクラスを使わない書き方で対処する方針で。

• jitclass

Numbaを使うと、エラーのデバッグに若干手間取ることがある。

特に、ローカルでは動くのにサーバではREが出た時に厄介。

Numbaはバージョンの違いで使える機能に結構差が出てくる。
AtCoderの現在のNumbaのバージョンは 0.48.0 なので、差し支えなければこれに揃える。
ただ、Numbaだけを揃えても、他のNumPy等のモジュールの違いのせいなのか不明だが、REが出ることがある。

• AtCoderでは、コンパイルフェーズにエラーが出ても、Pythonの場合はCEでなくREとなる
  • 参考: 提出 #16507701 - AtCoder Beginner Contest 177
  • コンパイルフェーズでRE、本実行でWAとなるはずのコードだが、結果はREとなっている

そのため、テストケースが全てREの場合はコンパイルフェーズの可能性も視野に入れる。

また、

• 「コードテスト」における実行の場合は、Pythonのコンパイルフェーズは行われないっぽい
  • 実行環境が、os.getcwd() == '/imojudge/sandbox' となるので、これによる切り分けが可能
• Numbaのエラーコメントはバックトレースが長く、コードテストの欄で最後まで表示されないことも多い

このため、

• AOTではそもそも実行できない
• JITに書き換えても、コンパイル時間が長いと10秒で打ち切られる
• 実行してエラーを発生させても、エラーコメントからエラー箇所が見えない

という問題により、ちょっとデバッグがしづらい。

これの解決策は今のところ見つけられていない。Numbaを選択するときのリスクの1つ。

Numbaの中身が複雑だったりしてコンパイルに時間がかかりすぎると、どうも勝手に打ち切られるらしい。

その場合、コンパイルフェイズでモジュールが生成されない→実行フェイズでモジュールが存在しないので、全てのテストケースで RE となる。

コンパイル時間を短くするような書き方については要調査。

• いもす研 (imos laboratory)
  • オンラインジャッジでコンパイルが終了しない場合は適切に打ち切る必要を語られている
• いもすさんはTwitterを使っています 「imos contestとか懐かしい．2008年とかもう8年前か………．使われていたのは第一世代 Imo Judge かしら？（※ AtCoderで使われているのは第三世代 Imo Judge」 / Twitter
  • 少なくとも2016年では上を元にしたジャッジシステムが使われているっぽい