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を使っていきましょう!