Python 競プロ用 Tips

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

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

• 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」という特殊な文字で判断するので、手元でのテストなど入力をコマンドプロンプト(Windows)から与える場合、そのままでは実行が始まらない
  • 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より使えるようになった。

競プロの文脈でポイントとなるのは、以下だろう。

• 高速化のため、NoPythonモード、事前コンパイル
• 使わない方がよい問題もある
  • 文字列の方が扱いやすい問題とか、多倍長整数が有効な問題とか
  • →全ての問題に共通して使うテンプレート、というのではなく、必要に応じて呼び出せるスニペットがよい
• 提出コードとローカル環境で同じコードで動かしたい

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

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

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、でよいと思う。

AtCoderのジャッジシステムでは、テストケース実行前に ONLINE_JUDGE を実行時オプションとした処理が1回走る。

• https://docs.google.com/spreadsheets/d/1PmsqufkF3wjKN6g1L0STS80yP4a6u-VdGiEv5uOHe0M/

C言語などにとってのコンパイル処理をするフェーズだが、 スクリプト言語にとっても起動を速くする効果があったりするため、行われている。

Numbaでのコンパイルも、このコンパイルフェーズを利用して行う。 コンパイルフェーズで生成されたファイルは削除されず残る。

ただ、ファイルが残る仕様は悪用しようと思えば色々できてしまうと思うので、もしかしたら将来的になくなるかも知れない。

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関数の中身だけ書けばいい」という状態にできる。

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

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

• それぞれを個別にコンパイル
• 1つの大枠の関数の中で個々の関数定義も書いて、大枠の関数のみコンパイル

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

普通に後者でいいと今のところは思っている。

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

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

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

• jitclass

以下を参照させていただいた。あまり変わっていないが、手元環境でも同一コードで動く。

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

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

基本的に 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)