Python で AtCoder を遊ぶときに知ってると便利かもしれないこと

気が向いたら後で検証したりで増やしたりするかも

割と当たり前のこととかそんなに使わないこととかが混じってる

PyPy を使う

これ必須
AtCoder で使える言語には PyPy があって Python の代わりに使える
Numpy が使えなかったりするけど速い
特にループが 3 重とかになるときに効果的
(オーバーヘッドが大きいので最初使ったときに微妙だなって勘違いしがち)

同じコードでの比較(もっと差がついたのもあった気がするけど忘れた)

Python (TLE) https://atcoder.jp/contests/abc073/submissions/2692814
PyPy (841ms) https://atcoder.jp/contests/abc073/submissions/2692842

入力の基本

下の 5 通りをコピペできるようにしておくと便利

# 整数 1 つ
N = int(input())

# 整数複数個
N, K = map(int, input().split())

# 整数 N 個 (改行区切り)
L = [int(input()) for i in range(N)]

# 整数 N 個 (スペース区切り)
A = list(map(int, input().split()))

# 整数 (縦 H 横 W の行列)
S = [list(map(int, input().split())) for i in range(H)]

input() を速くする

import sys
def input():
    return sys.stdin.readline()[:-1]

これを書いておくとよくわからないけど速くなる
特に PyPy だと効果的
というか PyPy の input() って Python より遅くないですか?なんで?

標準出力で便利なやつ

リストをスペース区切りで出力

arr = [1, 2, 3]
print(*arr)
1 2 3

リストを改行区切りで出力

arr = [1, 2, 3]
print(*arr, sep="\n")
1
2
3

最後に改行を入れない出力

print("あああ", end="")

処理を途中で終了するやつ

exit()

pow() の第3引数

pow(x, y, z) は x の y 乗を z で割った余りを返す
計算途中で巨大な数にならないので x**y % z より高速

mod = 10**9+7
print(pow(10, 10, mod))  # 10 の 10 乗を 1000000007 で割った余り
999999937

整数 x の mod 逆元は pow(x, mod-2, mod) で求められる
拡張ユークリッドの互除法を使うより速かった気がする

i 番目のアルファベットを整数 i に変換

ord()で ascii 的なやつに変換して 97 を引けばいい

"a" は 0 番目なので 0 になる

print(ord("a") - 97)
0

逆変換
1 番目のアルファベットは b なので "b" になる

print(chr(97 + 1))
b

再帰の深さの上限を増やす

何もしないと深い再帰で RE する

import sys
sys.setrecursionlimit(500000)

無限大

float("inf")

勝手にメモ化してくれる便利なやつ

デコレータ @lru_cache

from functools import lru_cache
@lru_cache(maxsize=None)  # メモ化したい関数の前につける
def fib(n):  # n 番目のフィボナッチ数を求める関数
    if n <= 1:
        return 1
    else:
        return fib(n - 2) + fib(n - 1)
print(fib(80))
37889062373143906

デフォルト値のある辞書型

collections.defaultdict がかなり便利
dd = defaultdict(lambda: 3) みたいにすると 3 が初期値の辞書ができる
グラフの問題とかいろんな問題で頻繁に使う

from collections import defaultdict
cnt = defaultdict(int)  # int を引数にすると 0 が初期値になる
# cnt = defaultdict(lambda: 0) と書いても良い

cnt[1] += 1
cnt[4] += 1
for i in range(5):
    print(cnt[i])
0
1
0
0
1

グラフのエッジを登録する例

from collections import defaultdict

# dafaultdict をデフォルト値とした defaultdict
E = defaultdict(lambda: defaultdict(lambda: float("inf")))

E[0][1] = 2  # ノード 0 から ノード 1 への長さ 2 のエッジを張る

itertools

組み合わせとか直積とか、普通なら再帰を使わないといけなくて面倒そうなやつらが一瞬で書ける

直積

from itertools import product
for t in product([1, 2, 3], ["a", "b"]):
    print(t)
(1, 'a')
(1, 'b')
(2, 'a')
(2, 'b')
(3, 'a')
(3, 'b')

直積その 2

from itertools import product
for t in product(["a", "b"], repeat=3):
    print("".join(t))
aaa
aab
aba
abb
baa
bab
bba
bbb

順列

from itertools import permutations
# ["a", "b", "c"] の並べ替えをすべて列挙する
for t in permutations(["a", "b", "c"]):
    print(t)
('a', 'b', 'c')
('a', 'c', 'b')
('b', 'a', 'c')
('b', 'c', 'a')
('c', 'a', 'b')
('c', 'b', 'a')

組み合わせ

from itertools import combinations
# ["a", "b", "c"] から 2 つ選ぶ組み合わせをすべて列挙する
for t in combinations(["a", "b", "c"], 2):
    print(t)
('a', 'b')
('a', 'c')
('b', 'c')

重複を許す組み合わせ

from itertools import combinations_with_replacement
# ["a", "b", "c"] から重複を許して 2 つ選ぶ組み合わせをすべて列挙する
for t in combinations_with_replacement(["a", "b", "c"], 2):
    print(t)
('a', 'a')
('a', 'b')
('a', 'c')
('b', 'b')
('b', 'c')
('c', 'c')

Python のビット演算は速くない

&1 をするより %2 をしたほうが速かったりする
PyPy だと互角くらい

ソートいろいろ

from operator import itemgetter
a = [(1, "c", 1), (1, "b", 3), (2, "a", 0), (1, "a", 2)]
print(sorted(a))  # 0 番目の要素でソート、先頭の要素が同じなら 1 番目以降の要素も見る
print(sorted(a, key=itemgetter(0)))  # 0 番目の要素だけでソート
print(sorted(a, key=itemgetter(0, 2)))  # 0 番目と 2 番目の要素でソート
print(sorted(a, key=lambda x: x[0] * x[2]))  # 0 番目の要素 * 2 番目の要素でソート
print(sorted(a, reverse=True))  # 降順にソート
a.sort()  # 破壊的にソート、sorted() よりも高速
print(a)
[(1, 'a', 2), (1, 'b', 3), (1, 'c', 1), (2, 'a', 0)]
[(1, 'c', 1), (1, 'b', 3), (1, 'a', 2), (2, 'a', 0)]
[(1, 'c', 1), (1, 'a', 2), (1, 'b', 3), (2, 'a', 0)]
[(2, 'a', 0), (1, 'c', 1), (1, 'a', 2), (1, 'b', 3)]
[(2, 'a', 0), (1, 'c', 1), (1, 'b', 3), (1, 'a', 2)]
[(1, 'a', 2), (1, 'b', 3), (1, 'c', 1), (2, 'a', 0)]