はじめに
- 競技プログラミングで使えるPython標準ライブラリを紹介します。
- 実際にPython標準ライブラリを使って解いたAtCoderの問題とコードも載せています。
- AtCoderでのPython3のバージョンは3.4.3です。
- 問題解説はしません。
標準ライブラリ
bisect
配列二分法アルゴリズム
bisect_left(a, x, lo=0, hi=len(a))
ソート済みのリスト$a$に対して、$x$を$a$に挿入できる位置を返します。$a$に$x$が含まれる場合、挿入箇所は既存のどの$x$の値よりも前（左）になります。
bisect_right(a, x, lo=0, hi=len(a))
bisect_left()とほぼ同じですが、$a$に$x$が含まれる場合、挿入箇所は既存のどの$x$の値よりも後（右）になります。
from bisect import bisect_left, bisect_right
a = [1, 2, 2, 2, 3, 5]
print(bisect_left(a, 2)) # -> 1
print(bisect_right(a, 2)) # -> 4
問題例
heapq
ヒープキューアルゴリズム/優先度キューアルゴリズム
itemをheapにpushします。
heapから最小の要素を返します。
from heapq import heappop, heappush
a = []
heappush(a, (10, 'x'))
heappush(a, (5, 'y'))
heappush(a, (1, 'z'))
print(a) # -> [(1, 'z'), (10, 'x'), (5, 'y')]
print(heappop(a)) # -> (1, 'z')
print(heappop(a)) # -> (5, 'y')
print(heappop(a)) # -> (10, 'x')
collections
コンテナデータ型
deque
両端でappendとpopを高速に行えるリスト風コンテナ。appendleft()やpoplift()は、listに対する同等の操作より高速に実現できます。
from collections import deque
q = deque(['a', 'b', 'b', 'b', 'd', 'd'])
print(q.popleft()) # -> a
print(q) # -> deque(['b', 'b', 'b', 'd', 'd'])
q.appendleft('z')
print(q) # -> deque(['z', 'b', 'b', 'b', 'd', 'd'])
ハッシュ可能なオブジェクトをカウントする辞書のサブクラス。
1要素の個数をカウントするには、listやtupleのcount()メソッドでもいいですが、全要素の個数をカウントするときに便利です。
from collections import Counter
l = ['a', 'b', 'b', 'b', 'd', 'd']
t = ('a', 'b', 'b', 'b', 'd', 'd',)
print(l.count('b')) # -> 3
print(t.count('b')) # -> 3
c = Counter(l)
print(c) # -> Counter({'b': 3, 'd': 2, 'a': 1})
# カウントが大きい順に全ての要素を返す
print(c.most_common()) # -> [('b', 3), ('d', 2), ('a', 1)]
itertools
効率的なループ実行のためのイテレータ生成関数
複数のiterableの直積を返します。
from itertools import product
l = ['a', 'b', 'c']
m = ['x', 'y']
print(list(product(l, m)))
# -> [('a', 'x'), ('a', 'y'), ('b', 'x'), ('b', 'y'), ('c', 'x'), ('c', 'y')]
iterableの要素から長さrの順列を返します。
from itertools import permutations
l = ['a', 'b', 'c']
print(list(permutations(l, 3)))
# -> [('a', 'b', 'c'), ('a', 'c', 'b'), ('b', 'a', 'c'), ('b', 'c', 'a'), ('c', 'a', 'b'), ('c', 'b', 'a')]
print(list(permutations(l, 2)))
# -> [('a', 'b'), ('a', 'c'), ('b', 'a'), ('b', 'c'), ('c', 'a'), ('c', 'b')]
combinations(iterable, r)
iterableの要素からr個選ぶときの組み合わせを返します。
from itertools import combinations
l = ['a', 'b', 'c', 'd', 'e']
print(list(combinations(l, 2)))
# -> [('a', 'b'), ('a', 'c'), ('a', 'd'), ('a', 'e'), ('b', 'c'), ('b', 'd'), ('b', 'e'), ('c', 'd'), ('c', 'e'), ('d', 'e')]
問題例
ABC123 A Five Antennas
コード
functools
高階関数と呼び出し可能オブジェクトの操作
lru_cache(maxsize, typed)
関数をメモ化用の呼び出し可能オブジェクトでラップするデコレータです。最近の呼び出しmaxsize回まで保存します。メモ化再帰で使えます。
from functools import lru_cache
@lru_cache(maxsize=None)
def fib(n):
if n < 2:
return n
return fib(n - 1) + fib(n - 2);
print(fib(100)) # -> 354224848179261915075
print(fib.cache_info()) # -> CacheInfo(hits=98, misses=101, maxsize=None, currsize=101)
reduce(function, iterable[, initializer])
2つ引数をとるfunctionにiterableの要素を左から累積的に適用していき、その結果を返します。
from functools import reduce
def f(a, b):
return a * b
l = [1, 2, 3, 4, 5]
# ((((1 * 2) * 3) * 4) * 5)と同値
print(reduce(f, l)) # -> 120
fractions
有理数
gcd(a, b)
整数$a$と$b$の最大公約数を返します。
最小公倍数を求めたいときは、以下のようにlcmを実装しましょう。
from fractions import gcd
def lcm(a, b):
return a*b // gcd(a, b)
※Python3.5で、math.gcd()が追加され、fractions.gcd()は非推奨になりました。
問題例
ABC070 C Multiple Clocks
コード
番外編
競プロで使える組み込み関数を紹介します。
pow()
累乗
$ pow(x, y[, z]) $ で、 $ x^y $ または $ x^y $ に対する $ z $ の剰余を返します。
これを使うことで、$ mod $ $ p $ における $ a $ の 逆元を、$ pow(a, p-2, p) $ で求められます。
$n$個から$k$個選ぶ組み合わせ数の$p$で割った余りを求めるとき等に使えます。
def nCk(n, k, p):
ret = 1
for i in range(k):
ret = ret * (n - i) * pow(i + 1, p - 2, p) % p
return ret
問題例
ABC034 C 経路
コード