競プロ的観点からPython3のドキュメントを読む

はじめに

Python3で競技プログラミングをするときに便利な仕様、標準ライブラリを探るべく、Python 3.8.2 ドキュメントを読みました。
これはあまり知られてなさそうですねと思ったり、普段他の言語を使っている人が python3 を使いたいときに困りそうですねと思ったものを列挙しました。
python3.8.2 のドキュメントのまとめなので、python3.8.2 未満では使えない機能が含まれています。実際に競プロに使うときはご注意ください。

• 引数リストのアンパック
• ループのテクニック
• べき乗演算
• 行列乗算演算
• 複素数型
• オーバーラップする代入文
• 比較の連鎖
• 同一性の比較
• 所属検査演算
• 短絡評価
• frozenset
• 疑似乱数生成
• accumulate
• 順列と組み合わせ

##引数リストのアンパック
https://docs.python.org/ja/3.8/tutorial/controlflow.html#unpacking-argument-lists
シーケンス型の頭に*演算子をつけることで、要素を展開したうえで関数に渡すことができます。計算結果の標準出力時をはじめに広く活用の機会がありそうです。

ans = [0, 1, 2]
print(*ans) # 引数をアンパックして関数に渡す、print(0, 1, 2)と等価

# 実行結果
0 1 2

##ループのテクニック
https://docs.python.org/ja/3.8/tutorial/datastructures.html#looping-techniques
辞書型の key と value、シーケンス型の index と value など、複数の値を並列に取り出すことができます。配列参照演算子[]によるアクセスが減り定数倍の高速化が期待できるうえ¹、添え字がすっきりするのでうれしいです。

# enumerate()でシーケンス型のindexとvalueを同時に取り出せる
for index, val in enumerate(["a", "b", "c"]): 
    print(index, val)

# 実行結果
0 a
1 b
2 c

# 辞書型のitems()メソッドを使うとkeyとvalueを同時に取り出せる
dictionary = {"January": 1, "February": 2, "March": 3}
for key, val in dictionary.items(): 
    print(key, val)

# 実行結果
January 1
February 2
March 3

# zip()を使うと複数のシーケンスを並列にループできる
a = [1, 3, 5]
b = [2, 4, 6]
for a_val, b_val in zip(a, b): 
    print(a_val, b_val)

# 実行結果
1 2
3 4
5 6

# 要素が少ない方に合わせてループが止まる。例えば隣接要素の差が求めらる
a = [1, 2, 4]
a_diff = []
for a0, a1 in zip(a, a[1:]): 
    a_diff.append(a1 - a0)
print(a_diff)

# 実行結果
[1 2]

##べき乗演算
https://docs.python.org/ja/3.8/reference/expressions.html#the-power-operator
標準でべき乗演算子**が用意されています。円の面積を求めたり、二点間の距離を求めたりするちょっとした幾何計算で活躍してくれることでしょう。²
組み込み関数の pow(base, exp) と等価ですが、pow には加えて mod p 上でのべき乗を計算する機能があります。大きな通り数になる数え上げ問題などで大活躍します。³

print(2**10)
print(1.5**-2.0) # base, expが整数でなくともよい

# 実行結果
1024
0.4444444444444444

print(pow(2, 10, 100)) # 2**10 mod 100

# 実行結果
24

##行列乗算演算
https://docs.python.org/ja/3.8/reference/expressions.html#binary-arithmetic-operations
標準演算子として行列乗算演算子@が用意されています。ただ、python3.8 時点で組み込み型に行列型はありませんので、数値計算モジュール numpy で用意されてる array型や matrix型で行列積を求めるときに使うとよいでしょう。
例えば行列の累乗を必要とする問題に使えそうですね。⁴

# 組み込み型に@演算子の実装を持つものがないため、numpy.arrayを使用
import numpy
a = numpy.array([[1, 2], 
                 [2, 1]])
b = numpy.array([[3, 4], 
                 [4, 3]])
print(a @ b) # numpy.dot(a, b)と等価

# 実行結果
[[11 10]
 [10 11]]

##複素数型
https://docs.python.org/ja/3.8/library/stdtypes.html#numeric-types-int-float-complex
標準で複素数型が用意されています。心置きなくFast Fourier Transformationが書けます。⁵

c = 1 + 1j # 虚数リテラルjを使って複素数型を表す
print(c)
print(c.conjugate()) # 共役複素数も参照できる
import math
print(pow(math.e, math.pi * 1.0j).real) # .realで実数部を、.imagで虚数部を取り出せる

# 実行結果
(1+1j)
(1-1j)
-1.0

##オーバーラップする代入文
https://docs.python.org/ja/3.8/reference/simple_stmts.html#assignment-statements
代入文で値の swap ができます。要素が参照しているオブジェクト間に依存関係がない限り、同時に評価されます。

a = 0; b = 1
a, b = b, a # 左辺と右辺がオーバーラップする代入は基本的に同時に実行され、swapとして使うことができる
print(a, b)

# 実行結果
1 0

i = 0
x = [0, 0]
i, x[i] = 1, 2 # 参照するオブジェクトに依存関係があるときは、左から右へ順に評価される 
print(x)

# 実行結果
0 2

##比較の連鎖
https://docs.python.org/ja/3.8/reference/expressions.html#comparisons
比較演算子を連鎖させることができます。範囲内に値が含まれるかどうか確認したいときに、同じ変数を何度も書かなくて済むのはありがたいですね。

a = 5
print(1 <= a < 7) # 比較演算は連鎖することができる、1 <= a and a < 7とほぼ等価だが、この書き方では短絡評価は行われない

# 実行結果
True

##同一性の比較
https://docs.python.org/ja/3.8/reference/expressions.html#is-not
==演算子は等価であるか判定します。
is演算子はオブジェクトの identity が同一かどうか判定します。
間違えて使うと解決が難しいバグを生みそうです、気をつけましょう。

a = [1, 2, 3]
b = a
print("a is b:", a is b) # id(a) == id(b)と等価
print("a == b:", a == b)

# 実行結果
a is b: True
a == b: True

a = [1, 2, 3]
b = a.copy()
print("a is b:", a is b) # id(a) == id(b)と等価、list.copy()は新しくオブジェクトを生成する
print("a == b:", a == b)

# 実行結果
a is b: False
a == b: True

##所属検査演算
https://docs.python.org/ja/3.8/reference/expressions.html#membership-test-operations
ある要素が list, tuple, set, dict などに存在するかどうかを評価します。set, dict などはO(1)で評価が完了するのに対し、list, tuple などでは O(n) の計算量となることに注意です。

a = {1, 2, 3}
print(2 in a)

# 実行結果
True

##短絡評価
https://docs.python.org/ja/3.8/library/stdtypes.html#boolean-operations-and-or-not
真偽値や関数を and や or で繋げた場合、評価が左から行われ、値が確定した時点でそれ以降の真偽が評価されなくなります。例えば配列の参照を行う前に、インデックスが配列長に収まるか確認することで配列外参照を防げそうです。

def check():
    print("checked!")
    return True

print(False or check()) # 左から右へと順にFalse, check()が評価される
print(True or check()) # 第一項で返り値がTrueと確定したため評価が停止され、check()は評価されない

# 実行結果
checked!
True

True

a = [1, 2, 3]
# 第一項でインデックスが配列長に収まることを確認しておけば、第二項の配列参照エラーを未然に防ぐことができる
if 2 < len(a) and a[2] == 3:
    print("a[2] == 3")

# 実行結果
a[2] == 3 

##frozenset
https://docs.python.org/ja/3.8/library/stdtypes.html#set-types-set-frozenset
set はミュータブルな型でありhash値を計算できません。frozenset はイミュータブルである代わりにhash値を計算できる集合型です。
無向辺の辞書を作ったりできそうですが、そもそも集合型はあまり速くないので慎重に使いましょう。

dictionary = dict()
dictionary[frozenset((1, 2))] = 3 # frozensetはimmutableでhash化できるので、dictのkeyとして使えます
dictionary[frozenset((2, 3))] = 5
print(dictionary[frozenset((3, 2))]) # 集合そのものがkeyなので、要素の順番は問わない

# 実行結果
5

##疑似乱数生成
https://docs.python.org/ja/3.8/library/random.html#module-random
randomモジュールはメルセンヌツイスタを乱数生成器に、周期 2**19937-1 の疑似乱数を作ります。例えば randrange で指定範囲のランダムな整数の生成が、shuffle でランダムな順列が生成できるのでランダムなテストケースが作れてうれしいことがありそうです。

import random
a = [random.randrange(1, 7) for _ in range(10)] # start <= n < stopであるような整数nを返す
print(a)

a = [1, 2, 3, 4, 5, 6]
random.shuffle(a) # 可変なシーケンス型をインプレースにシャッフルする
print(a)

# 実行結果
[3, 6, 4, 5, 3, 6, 6, 2, 3, 2]
[4, 2, 5, 6, 3, 1]

##accumulate
https://docs.python.org/ja/3.8/library/itertools.html#itertools.accumulate
累積的な演算の結果をイテレータ型で返します。累積和のみならず様々な累積の演算をワンライナーで実現できてしまいます。便利ですね。

from itertools import accumulate
a = [1, 3, 2, 5, 4]

# 累積和がワンライナーで書ける
print(list(accumulate(a, initial=0)))

# 実行結果
[0, 1, 4, 6, 11, 15]

# 組み込み関数を渡すことで、和以外の累積演算も書ける
print(list(accumulate(a, max, initial=0))) # max関数を渡すことで累積maxを得る

# 実行結果
[0, 1, 3, 3, 5, 5]

# ラムダ式を渡すことで、ニッチな累積演算をワンライナーで書ける
func = lambda x, y : x * y % 10
print(list(accumulate(a, func, initial=1))) # 累積の積を10で割った余り

# 実行結果
[1, 1, 3, 6, 0, 0]

##順列と組み合わせ
https://docs.python.org/ja/3.8/library/math.html#number-theoretic-and-representation-functions
標準で順列の場合の数を返す関数 perm と組み合わせの場合の数を計算する関数 comb が用意されています。n が大きいときは返り値が指数的に大きくなり、演算が遅くなる恐れがあるので気をつけましょう。

import math
print(math.perm(5, 2)) # 5C2を計算します
print(math.comb(5, 2)) # 5P2を計算します

# 実行結果
20
10

おわりに

他にもみなさんが、この仕様、標準ライブラリはあまり知られてなさそうですねと思ったり、普段他の言語を使っている人が python3 を使いたいときに困りそうですねと思ったものがあれば、教えていただけると幸いです。

参考

Python 3.8.2 ドキュメント

1. 配列参照演算子の方が約4倍ほど遅いです。Pythonの知っておくと良い細かい処理速度の違い8個：https://www.kumilog.net/entry/python-speed-comp ↩

2. 実はn次元空間の二点間のユークリッド距離を求める組み込み関数があります。math.dist()です。 ↩

3. 「1000000007 で割ったあまり」の求め方を総特集！ 〜 逆元から離散対数まで 〜: https://qiita.com/drken/items/3b4fdf0a78e7a138cd9a ↩

4. NumPyの行列クラスnumpy.matrixには行列のべき乗が**演算子に実装されており、それを使ってしまえばよいという話もあります。 ↩

5. NumPyにはnumpy.fft.fftというFFTの実装もあります。https://docs.scipy.org/doc/numpy/reference/generated/numpy.fft.fft.html ↩