どういう記事

どうも。私(@rainline00)がAtcoder Begginers Contestの過去問をpythonで解いていく過程で必要になった競プロ独特の書き方をまとめていきます。随時追記していきます。

index

書き方

入力

n行に１つずつ入力される整数が地味に忘れやすいです。

s = input() #文字列
n = int(input()) #整数値
a, b = map(int, input().split()) #2つ以上の整数
li = list(map(int, input().split())) #list
li2 = [int(input()) for _ in range(n)] #n行に１つずつ入力される整数

リストの要素検索

list.index(n)では一致した値のインデックスを一つしか返さない。リスト内包表記を用いて複数の値をリストにまとめて返す。

li = [0, 2, 1, 4, 2, 3, 2] 
indexes = [i for i, x in enumerate(li) if x == 2] # 2を検索
print(indexes) # [1, 4, 6]

全探索

bit 全探索

各変数について「使うor使わない」のような場合を全て列挙しても間に合う場合（$O(2^N)$で間に合う場合）、bit 全探索が使えます。「使うor使わない」を表した全体の状態を、変数の数だけ桁数をとった
２進数1つだけで表すことができます。

for state in range(1 << N):
    # リスト内包表記でその状態の各要素のフラグを作ることができる
    isOn = [True if (state >> i) % 2 == 1 else False for i in range(N)]

例題: ABC167C - Skill Up

各参考書を「買うor買わない」の$2^n$個の状態を全て列挙し、条件を満たしているかどうか判定します。numpyのarrayを用いることで、高速かつ簡潔に記述することができます。

abc167.py

import numpy as np

n, m, x = map(int, input().split())
li = [np.array(map(int, input().split())) for _ in range(n)]
ans = 10 ** 10
comp = np.array([x for _ in range(m)])

for i in range(1 << n):
    sum = np.array([0 for _ in range(m + 1)])
    now = i
    for j in range(n):
        di = now % 2
        now = now >> 1
        if di == 1:
            sum += li[j]

    if all(sum[1:m+1] >= comp):
        ans = min(ans, sum[0])

if ans == 10 ** 10:
    print(-1)
else:
    print(ans)

最大公約数・最小公倍数

多変数の最大公約数・最小公倍数を求める場合は、高階関数reduce()を使って再帰的に書くことができる。
~~atcoderのpythonはversion3.4.3(2019/12/31現在)なのでmathモジュールではなくfractionsモジュールを使う。~~
（2020/5/19 追記）AtcoderのPythonが3.8.2にアップデートされました！ mathモジュール使いましょう

import fractions
from functools import reduce

def gcd_list(numbers):
    return reduce(fractions.gcd, numbers)

def lcm(x, y):
    return (x * y) // fractions.gcd(x, y)

def lcm_list(numbers):
    return reduce(lcm, numbers)

累積和

ある区間の和を$O(1)$で求めるアルゴリズム。前処理に$O(N)$かかる。
数列$\{a_n\}$に対して$s_0=0, s_i=\sum_{k=0}^{i-1}a_k$を満たす数列$\{s_n\}$を定義する。
つまり、$s_i$は$[0, i)$の和である。
$[a, b)$の和が$s_b-s_a$で求められる。

a = [i for i in range(11)]
s = [0]
for i in a:
    s.append(s[-1] + i)
print(s[3] - s[0]) # 0 + 1 + 2 = 3
print(s[11] - s[8]) # 8 + 9 + 10 = 27

組み合わせ・順列

総数を出す

AtcoderのPythonはVersion3.4.3(2019/12/31時点)なので、`scipy`モジュールの`scipy.special.perm()`は使えないが、`scipy.misc.comb()`は使える。高速なのでscipy使いたいんだけどな。

（2020/5/19 追記）AtcoderのPythonが3.8.2にアップデートされました！scipy.special使いましょう！

from scipy.misc import comb
print(comb(3, 2)) # 3.0
print(comb(3, 2, exact=True)) # 3

順列は自前で実装する。

from math import factorial
def perm(n, r):
    return factorial(n) // factorial(n - r)
print(perm(3, 2)) # 6

深さ優先探索・幅優先探索（編集中）

どちらも可能な状態を列挙して探索する。探索の仕方が違う。
pyてょん日記さんの競プロ覚書：深さ優先探索，幅優先探索まとめを参考にさせていただきました。

深さ優先探索

根付き木をイメージする。根から末端の葉まで一直線に探索する。

初期状態から始める
行けるところまで一直線で遷移する
2.で行けるところまで行ったら一つ遷移を戻してそこから行けるところまで遷移していく
2.と3.を繰り返し状態を列挙する

以上の工程を実装する。重要なのは

終了条件を明確にする
現在の状態と深さを対応つけて考える


def dfs（）:

幅優先探索

キューの実装にcollectionsモジュールのdequeを使う。 listを使ってキューを再現することもできるが、pop(0)やinsert(1, n)の計算量が$O(n)$かかる。対して、dequeならばそれらの計算が$O(1)$でできる。必ずこちらを使おう。

Pythonで解く競プロテクニック集