読者です 読者をやめる 読者になる 読者になる

n-knuu's logs

憧れ駆動。だいたい競プロ

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モジュール内の関数なので注意。

終わりに

皆さんもC++のサブウェポンとしてpython使いましょう!

明日は @bakaming さんです。

*1:まだ11日だからきっとセーフ

*2:考えていたネタで上手く記事が書けなかった