pythonで競技プログラミングをする際に便利な関数とか5選
この記事は Competitive Programming Advent Calendar 2015 の11日目の記事です。*1
昨日は @yuizumi_y5iさんでABC を bash で解いた話 - yuizumi’s diaryでした。
この記事は？
突然ですが、競技プログラミングをする上で便利だ思うpythonの関数/メソッドを(大きく分けて)5つ紹介します。*2
この記事のpythonのバージョンは3.4.3を想定しています。
itertools系
itertools.permutationsが階乗、itertools.productが重複順列、itertools.combinationsが組合せ、itertools.combinations_with_replacementが重複組合せ
>>> list(itertools.permutations([1, 2, 3])) # 3! 通り [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)] >>> list(itertools.product([1, 2, 3], repeat=2)) # 3**2 通り [(1, 1), (1, 2), (1, 3), (2, 1), (2, 2), (2, 3), (3, 1), (3, 2), (3, 3)] >>> list(itertools.combinations([1, 2, 3, 4], 2)) # 4C2 通り [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)] >>> list(itertools.combinations_with_replacement([1, 2, 3], 3)) # 3H3個 [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 2, 2), (1, 2, 3), (1, 3, 3), (2, 2, 2), (2, 2, 3), (2, 3, 3), (3, 3, 3)]
C++だと階乗はnext_permutationを使えば良いが、重複順列は再帰関数で書いたり、組合せは蟻本中級編pp.143~145のようなビット演算で求めたりする必要があるが、pythonだとライブラリ一発で書けるので便利。
collections.Counter
リスト/文字列の要素数をカウントする
>>> import collections >>> collections.Counter('abccbabcabcbaba') Counter({'b': 6, 'a': 5, 'c': 4})
Counter同士で四則演算もできることが役に立つこともある(参考: Submission #488618 - CODE FESTIVAL 2014 予選B | AtCoderとかSubmission #11571960 - Codeforces)
また、most_commonを使うと簡単にカウント数最大/最小の要素が取り出せる。
>>> c = collections.Counter('abccbabcabcbaba') >>> mc = c.most_common() >>> mc [('b', 6), ('a', 5), ('c', 4)] >>> mc[0] # 最大 ('b', 6) >>> mc[-1] # 最小 ('c', 4)
8.3. collections — コンテナデータ型 — Python 3.4.3 ドキュメントを見るだけでも色々使い方が分かる
print関数
print関数は便利で、リストの要素を空白ごとに表示する際も、一行ずつ表示する際にも使える。
>>> a = [1, 2, 3, 4, 5] >>> print(*a) 1 2 3 4 5 >>> print(*a, sep='\n') 1 2 3 4 5
これを応用すれば、2次元リストも以下で表示できる。
>>> matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] >>> [print(*row) for row in matrix] 1 2 3 4 5 6 7 8 9 [None, None, None]
[None, None, None]は標準出力には表示されないので問題ない。
str.format
桁数指定、基数変換とかいろいろ出来る。
>>> '{:.6}'.format(3.14159265) '3.14159' >>> '{0:b} {0:o} {0:x}'.format(10) '1010 12 a'
単に基数変換ならformat関数とint関数でもできる。(参考: Pythonで競技プログラミングする時に知っておきたいtips(データ構造編) - Qiita)
str.maketrans, str.translate
str.maketransで辞書を作って、str.translateで文字列を変換、ということをやってくれる。
>>> s = 'abccba' >>> s.translate(str.maketrans({'a': '1', 'b': '2', 'c': '3'})) '123321' >>> s.translate(str.maketrans('abc', '123')) # これは上と同じ '123321'
例えば A: お買い物クライシス - AtCoder Regular Contest 019 | AtCoder だと、
Submission #144862 - AtCoder Regular Contest 019 | AtCoder のように、replaceを数珠つなぎする解法に対して、str.maketrans/str.translateを使うと、
Submission #378402 - AtCoder Regular Contest 019 | AtCoderとすっきり書ける。
maketransは、python2系ではstringモジュール内の関数なので注意。
