Pythonで競技プログラミング -ライブラリ編-

okumuraです。前回の記事で異様にいいねがついて少々驚きました。その記事の最後に「余力があればよく使うライブラリー集とかも出すかもしれません」とかいってて何もしてなかったので、まとめました。

0. はじめに

この記事は半分自分用のために作りました。私自身あんまりライブラリを使わない方なので、結構抜けているかもしれません。「推しライブラリがないぞ！」などなどはコメントに記入していただけると幸いです。また説明が雑な箇所や個人的見解、ミスなど多々見られるかもしれません。これは私の理解力不足と経験不足などなどからくるものです、申し訳ございません。

一応各項目の後には私のAtCoder,Codeforcesでの使用例も載せておきました。問題のネタバレも含むので読むときには注意してください。参考になればいいと思います。

1. sys

sys.exit

プログラムを終了させます。私がsysモジュールでよく使うのはこれです。自明な条件分岐をすぐに終わらせれるから非常に楽です。例えば「N=1の時はYESでそのほかの時は・・・」というプログラム書きたい時はこんな使い方とかよくします。

import sys
N=int(input())
if N==1:
    print('YES')
    sys.exit()

結構簡単に条件分岐を減らせて、ネストが深くなるのを防げるので、かなり使っています。
AGC028-A Two Abbreviations 解答例

(追加事項) shiracamusさんから「exit(), quit()でもいいのでは？」というご指摘がありました。競プロではそちらでもいいと思います。

sys.stdin

入力がめっちゃ早くなります。ただ私はあんまり使っていないです。今すぐできる簡単な使い方は整数値などの入力に対して

import sys
input = sys.stdin.readline

とやってあとはいつも通り使うかんじです。ただし文字列を扱う時には注意すること。改行やらなんやらで事故ります。（事故った例を見たことあるので、文字列には使わない方が良い）。詳しくはPythonで競プロやるときによく書くコードをまとめてみたを参照してください。
CodeForces E. Polycarp's New Job 解答例

sys.setrecursionlimit

Pythonでは、再帰呼び出しの上限がデフォルトで1000回らしくこの上限を上げます。「Pythonで再帰関数のプログラム書いて提出したのにREになったんだが・・・」って時の解決法です。再帰関数使う時にはこれをつけておくとよいです。

import sys
sys.setrecursionlimit(10**6)

ARC031-B 埋め立て 解答例
EDPC-P Independent Set 解答例

2. collections

collections.deque

デックが使えるようになります。デックは前から後ろから出しても入れても計算量$O(1)$のデータ構造です。幅優先検索、深さ優先検索、トポロジカルソート、01BFSなどでよく使います。

l = [0,1,2,3]
q = deque(l)
q.append(4) # 後ろから4を挿入, l=deque([0,1,2,3,4])
q.appendleft(5)#前から5を挿入, l=deque([5,0,1,2,3,4])
x = q.pop() #後ろの要素を取り出す, x=4, l=deque([5,0,1,2,3])
y = q.popleft() # 前の要素を取り出す, y=5, l = deque([0,1,2,3])

ABC007-C 幅優先検索　解答例

collections.defaultdict

なんか好きなのでよく使います。普通の辞書と違って存在しないキーに参照してもエラーになりません。詳しくはPythonでdictからデフォルト値指定をしつつ値を取得する参照にしてください。個数カウントする時に下のようによく使います。

from collections import defaultdict
A=[int(i) for i in input().split()]
dd=defaultdict(int)
for a in A:
    dd[a]+=1

AGC031-A Colorful Subsequence 解答例

あとは個人的に辞書で動的計画法する時のもよく使いますが、これは真似しなくても良いです。
ARC067-E Grouping 解答例

collections.Counter

個数カウントするならこっちを使っている人の方が多いと思います。私はあんまり使ってないですが・・・。詳しくはPythonのCounterでリストの各要素の出現個数をカウント参照にしてください。辞書っぽい扱いができるのと「個数が多いものを２取る」といったことができるのが良い点だと思います。

from collections import Counter
l=['a','b','b','c','b','a','c','c','b','c','b','a']
S=Counter(l)#カウンタークラスが作られる。S=Counter({'b': 5, 'c': 4, 'a': 3})
print(S.most_common(2)) #[('b', 5), ('c', 4)]
print(S.keys()) #dict_keys(['a', 'b', 'c'])
print(S.values()) #dict_values([3, 5, 4])
print(S.items()) #dict_items([('a', 3), ('b', 5), ('c', 4)])

ARC103-C /\/\/\/ 解答例

3. itertools

itertools.accumulate

累積和が一発でできます。ものすごく速いため便利です。ですが私はあんまり使ってないです（なんか初期値で混乱するのとtypeがわからなくなるため）。

from itertools import accumulate
A=[1,4,3,4,6,5]
print(list(accumulate(A))) #[1, 5, 8, 12, 18, 23]

先ほどのCounterとaccumulateを使うと1行解答ができます。
AGC023-A Zero-Sum Ranges 解答例

itertools.permutations, itertools.combinations

私はほぼ使ったことないです(他にもproductというものがあるが、使ってすらいなかった)。しかもそこまで速くならないらしいです。Lをリスト（または文字列）, rを自然数として
- permutations(L,r) Lのr個の順列を全て列挙
- combinations(L,r) Lのr個の組み合わせを全て列挙
をそれぞれ出します。公式ドキュメントを見たら具体例ものっててわかりやすいです。

from itertools import product, permutations,combinations
A=[1,2,3,4]
for i in permutations(A,2):
    print(i,end=' ')
#(1, 2) (1, 3) (1, 4) (2, 1) (2, 3) (2, 4) (3, 1) (3, 2) (3, 4) (4, 1) (4, 2) (4, 3) 
for i in combinations(A,2):
    print(i, end=' ')
#(1, 2) (1, 3) (1, 4) (2, 3) (2, 4) (3, 4) 

• permutation ABC054-C One-stroke path 解答例
• combinations Tenka1 Programmer Contest-D Crossing 解答例

4. operator

operator.itemgetter

座標などどのキーでソートするかを指定できる。例見た方が早い。

from operator import itemgetter
B =[(5,8), (6,10), (7,2),(4,1), (3,11),(9,0)]
print(sorted(B, key = itemgetter(0))) #第1変数で昇順ソートしてる
#[(3, 11), (4, 1), (5, 8), (6, 10), (7, 2), (9, 0)]
print(sorted(B, key = itemgetter(0),reverse=True)) #第1変数で降順ソートしてる
#[(9, 0), (7, 2), (6, 10), (5, 8), (4, 1), (3, 11)]
print(sorted(B, key = itemgetter(1))) #第2変数で昇順ソートしてる
#[(9, 0), (4, 1), (7, 2), (5, 8), (6, 10), (3, 11)]
print(sorted(B, key = itemgetter(1),reverse=True)) #第2変数で降順ソートしてる
#[(3, 11), (6, 10), (5, 8), (7, 2), (4, 1), (9, 0)]

ABC091-C 2D Plane 2N Points 解答例

5. bisect

bisect.bisect_left, bisect.bisect_right, bisect.bisect

二分探索(にぶたん)でよく使います。ソートされたリストLに対して用いいます。毎回どっちがどっちって悩むのですが、

• bisect_left(L,x) xをLに挿入できる点(の番号)を探し当てる、ただしLにxがある場合は一番左になるようにする。
• bisect_right(L,x), bisect(L,x) xをLに挿入できる点(の番号)を探し当てる、ただしLにxがある場合は一番右になるようにする。

という違いがあります。詳しくは公式ドキュメントも参照してください。自分でもぴったりくる説明がピンとないので毎回困ります。計算量は$O(\log N)$ (NはリストLの要素数)です。

from bisect import bisect_left,bisect
L=[1,3,6,6,6,9,11]

print(bisect_left(L,0)) #0 
#0はLの先頭にくるので0番目に入ります、よって0が出力される。
print(bisect(L,0)) #0 上に同じ

print(bisect_left(L,5)) #2 
#5は3の次に入るので2番目に入ります。
print(bisect(L,5)) #2 上に同じ

print(bisect_left(L,6)) #2
#同じものがある場合は一番左の番号を出力するので、3の次にこの6は入ります。よって2番目に6がはいる。 
print(bisect(L,6)) #5 
#同じものがある場合は一番左の番号を出力するので、6の次にこの6は入ります。よって5番目に6がはいる。

print(bisect_left(L,12)) #7 
#12は一番後ろにくるので7番目に入ります。
print(bisect(L,12)) #7 上に同じ

0は自然数なのででてくる数字も0-indexです（これがさらに混乱を生み出す）。

codeFlyer予選-C - 徒歩圏内 解答例
ARC084-C Snuke Festival 解答例

6. heapq

heapq.heappop, heapq.heappush

プライオリティーキューのことです(Pythonだとヒープキューっていうらしい)。プライオリティーキューとは
- 計算量$O(\log N)$で要素を挿入
- 計算量$O(\log N)$で最小値を取り出す
ができるデータ構造です。最短経路を求めるダイクストラ法でお世話になります。

from heapq import heappop,heappush
L=[3,0,2,5,7,2]
H=[]
for l in L: #ここはH=heapq.heapify(L)でもいいです。
    heappush(H,l) 
print(H) #[0, 3, 2, 5, 7, 2]
print(heappop(H)) #0
print(heappop(H)) #2
print(heappop(H)) #2

• プライオリティーキューの使用例 ABC116-D Various Sushi 解答例
• ダイクストラ法 SoundHound Inc. Programming Contest 2018 予選-D Saving Snuuk 解答例

7. fractions

fractions.gcd

最大公約数を出してくれます。自分でユークリッドの互除法作らなくてもいいです。ただしAtCoderのPythonのバージョンは3.4.3です。よってgcdはfractionsモジュールの中にあります。誰しもが一回はハマるので注意しましょう(Python 3.5からmathモジュールに追加されました)。

from fractions import gcd
print(gcd(78627872,1798742872)) #8

ABC118-C Monsters Battle Royale 解答例

8. math

math.ceil, math.floor

• ceil 数字の切り上げ
• floor 数字の切り下げ

from math import ceil,floor
print(7/3) #2.3333333333333335 勝手にfloatにしてしまいます。
print(ceil(7/3)) #3
print(floor(7/3)) #2
print(7//3) #2 int型が出てきます

math.sqrt

正の平方根を出してくれます。自然数$A$の素因数分解を$O(\sqrt{A})$でやるときに使うことがあります。

CADDi 2018-C Product and GCD 解答例

math.cos, math.sin, math.pi

cos, sin, $\pi$が使えるようになります。一回CodeForcesでお世話になりました。

Codeforces-C NN and the Optical Illusion 解答例

9. copy

copy.deepcopy

A=[1,2,3]
B=A
B[1]=100
print(A) #[1, 100, 3]

皆さんもこのような経験があると思います（私はこれで３時間潰したことがあります）。これを防ぐのに使うのがcopy.deepcopyです。

from copy import deepcopy
A=[1,2,3]
B=deepcopy(A)
B[1]=100
print(A)

またcopy.copyというものもあります。copy.copyとcopy.deepcopyの違いに関してはPythonのcopyとdeepcopyについてを参照にしてください。２次元配列で違いが出てきますが、めんどくさいのでdeepcopy使えば良いと思います。

ARC083-D Restoring Road Network 解答例

10. numpy, scipy, scikits

AtCoderでは使えます、ちゃんとルールにも書いてます。でも競技プログラミングではあんまり使わないと思われます（機械学習とか深層学習では絶対使うんだけども）。いや無理すれば使えるのですが、別に使わなくてもなあ・・・って感じのが多いです。

一応何個かの応用例を書いておきます。
- numpy 行列演算ができる。ABC113-D Number of Amidakuji 解答例
- scipy.optimize.newton ニュートン法が使える。ABC026-D 高橋君ボール1号 解答例
他にもscipyでワーシャルフロイド法が使えるらしいです。

おまけ. 自分で作らなくてもPythonに組み込まれているもの

ユークリッドの互除法を作らなくてもfractions.gcdで最大公約数を求められました。実は同じような奴がもう一個あります

繰り返し二乗法

「$n^m$ を$p$で割った余り」を$O(\log m)$で求めるアルゴリズム。pythonだと

pow(n,m,p)

で求められます(自分で作らなくて良い！)。なので$p$が素数の時、$n$の$modulo$ $p$での逆元はフェルマーの小定理より

pow(n,p-2,p)

で求められます。計算量は$O(\log p)$です。二項係数とかでよく使います。

エクサウィザーズ 2019-E Black or White 解答例

-1. さいごに

自分用のメモの整理のために作りました。案外自分のためになりました。特にitemgetterやCounterは毎回使い方を調べてたのでこれを機に覚えれたかもしれません。