roiti46's blog

主に競プロ問題の解説を載せてます

pythonで競技プログラミングをする際に便利なほかの関数とか5選

この記事はCompetitive Programming (その2) Advent Calendar 2015の21日目の記事です(遅れてすみません)

20日目はMさんの高速なビット行列乗算
22日目はir5さんのtopcoder ボツ問供養です。

なにを書くのか?

Competitive Programming Advent Calendar 2015で11日目のn_knuu6さんがpythonで競技プログラミングをする際に便利な関数とか5選を紹介されてました。
今回はその記事では紹介されていなかった競プロに便利な関数・データ構造をいくつか紹介したいと思います。

sys.recursionlimit

pythonで競プロを使い始めた人がよくハマる落とし穴として再帰関数の再帰深さ制限があります。

>>> def mysum(n):
...     if n == 1: return 1
...     return n + mysum(n - 1)
... 
>>> mysum(999)
499500
>>> mysum(1000) #デフォルトの再帰深さ制限を超えてしまう
(省略)
RuntimeError: maximum recursion depth exceeded

デフォルトだと1000なので、それ以上の再帰深さがありそうなときはとりあえず大きめな値に設定しておきましょう。

>>> import sys
>>> sys.setrecursionlimit(100000)
>>> mysum(1000)
500500
>>> mysum(10000)
50005000

copy.deepcopy

pythonで1次元配列をコピーしたい場合は次のように行います。[:]は参照ではなく実際の値を返します。

>>> a = [1, 2, 3, 4, 5]
>>> b = a[:] #ここでb=aとするとaとbは同じ配列を指してしまう
>>> b
[1, 2, 3, 4, 5]
>>> b[0] = -1
>>> a
[1, 2, 3, 4, 5]
>>> b
[-1, 2, 3, 4, 5]

しかし2次元以上の配列ではうまくいきません。

>>> a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
>>> b = a[:]
>>> b
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
>>> b[0][0] = -1 #a[0][0]とb[0][0]は同じ要素を指している
>>> a
[[-1, 2, 3], [4, 5, 6], [7, 8, 9]]

2次元以上の配列をコピーしたいときはcopy.deepcopyを使いましょう。

>>> import copy
>>> a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
>>> b = copy.deepcopy(a)
>>> b
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
>>> b[0][0] = -1
>>> a
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
>>> b
[[-1, 2, 3], [4, 5, 6], [7, 8, 9]]

普通にコピーしてなぜダメなのかわからない人は、pythonのリストがどういう仕組みになっているのか一度見てみるといいかもしれません。

heapq

要素の追加、最小値のpopがO(logN)で行えるデータ構造です。

>>> import heapq
>>> a = []
>>> heapq.heappush(a, 5)
>>> heapq.heappush(a, 1)
>>> heapq.heappush(a, 3)
>>> a[0]
1
>>> heapq.heappop(a)
1
>>> a = [5, 10, 7, 4, 8, 1, 3]
>>> heapq.heapify(a) #破壊的変更
>>> a
[1, 4, 3, 10, 8, 7, 5]

C++のpriority_queueのようなものですが、pythonでは最小値が木の根にくるようになっています。
最大値バージョンが欲しい場合は値に-1をかけてpushしましょう。

>>> a = []
>>> heapq.heappush(a, -5)
>>> heapq.heappush(a, -3)
>>> heapq.heappush(a, -1)
>>> -heapq.heappop(a) #-1をかけるのを忘れない
5

defaultdict

collectionsライブラリの中の関数で、存在しないkeyについて参照しようとすると自動でvalueにデフォルト値を設定してkey-valueの組を作ってくれます。いちいち存在確認をしないでいいので自分もよく利用してます。

>>> from collections import defaultdict
>>> d = defaultdict(int) #intは引数無しだと0を返す関数
>>> d[0] = 10
>>> d[0]
10
>>> d[1] #存在しないkeyを参照しようとすると...
0

defaultdictの引数には値を返す関数を与えればいいので、lambda式を使えば任意の値をデフォルト値に設定することができます。

>>> d = defaultdict(lambda : 10000)
>>> d[1]
10000
>>> d = defaultdict(lambda : [1, 2, 3])
>>> d[1]
[1, 2, 3]

deque

これもcollectionsライブラリの中の関数です。名前の通りlistに対するpopとappendが両端で高速に行えます。pythonのlistはC++vectorのようなもので先頭でのpop,appendは低速です。したがって、BFSなどでlistの先頭からpopしていく場合で実行時間が厳しい時はdequeを使ってみるといいかもしれません。

>>> from collections import deque
>>> a = deque()
>>> a.append(1) #末尾に追加
>>> a.append(2)
>>> a.appendleft(3) #先頭に追加
>>> a
deque([3, 1, 2])
>>> a.pop() #末尾をpop
2
>>> a.popleft() #先頭をpop
3
>>> a
deque([1])

おまけ

余り知られていませんがpythonにはinfという値が存在します。infに数字やinfを足したりかけたりしてもinfです。ただし0をかけたり-infを足したりすると不定(nan)になります。配列の初期値として使える場面があるかも…。

>>> x = float("inf")
>>> x + 1
inf
>>> x * 10
inf
>>> x / 5
inf
>>> x * -1
-inf
>>> x * 0
nan
>>> x + (-x) #inf - inf
nan

おわりに

競プロにおいて速度は心もとないですが使いこなせると何かと便利だと思うので、ぜひpythonを使っていきましょう!