概要
Atcoderでよく用いるpythonのコードをすぐに参照するための自分用のチートシートです。
アルゴリズムについて、詳しく知りたい方は、他の記事をご参照ください。
演習をこなす度に、詰まった点を中心にアルゴリズム等を追記していく予定です。
基本操作
入力
sample.py
# 文字列の入力
s = input()
S = input().split()
# 整数、数列の入力
num = int(input())
N = list(map(int, input().split()))
最大公約数、最小公倍数
sample.py
# 最大公約数
def gcd(a, b):
while b:
a, b = b, a % b
return a
# 最小公倍数
def lcm(a, b):
return a * b // gcd (a, b)
素因数分解
sample.py
# nを素因数分解したリストを返す
def prime_factorize(n):
a = []
while n % 2 == 0:
a.append(2)
n //= 2
f = 3
while f * f <= n:
if n % f == 0:
a.append(f)
n //= f
else:
f += 2
if n != 1:
a.append(n)
return a
約数の列挙
sample.py
def divs_list(num):
divs = []
i = 1
while i*i <= num:
if num % i == 0:
divs.append(i)
if num // i != i:
divs.append(num//i)
i += 1
return divs
素数の列挙
sample.py
def sieve_of_eratosthenes(x):
nums = [i for i in range(x+1)]
root = int(pow(x, 0.5))
for i in range(2 ,root + 1):
if nums[i] != 0:
for j in range(i, x+1):
if i*j >= x+1:
break
nums[i*j] = 0
return set(nums)
組み合わせnCrの余りを含めた計算
sample.py
def comb(n, r, mod):
r = min(r, n-r)
a, b = 1, 1
for i in range(r):
a = a*(n-i)%mod
b = b*(i+1)%mod
return a * pow(b, mod-2, mod) % mod
小数をより正確に扱いたい時
sample.py
from decimal import Decimal
# xには事前に数字が定義されているとする
x = Decimal(x)
リスト
sample.py
l = [0, 1, 2, 3]
# i番目以降（iを含む）
l[i:]
# i番目まで(iを含まない)
l[:i]
# 多次元配列のソート
## 1つの値でソート
ll = sorted(data, key=lambda x:x[1])
## 複数キー
ll = sorted(data, key=lambda x:(x[1], x[2]))
アルゴリズム関連
二分探索
sample.py
# 入力を受け取る
N, K = map(int, input().split())
A = list(map(int, input().split()))
# 配列の右端をok, 左端をngと設定する。それぞれ端から１ずれているのは、両端が未判定のため。
ok = N
ng = -1
# while文で繰り返し真ん中の値を調べていく。while文を抜け出す条件は、okとngが隣り合う（差が1になる）まで
while abs(ok - ng) > 1:
# 真ん中の値をmidに代入
mid = (ok + ng) // 2
# 真ん中の値がK以上ならokにmidを、そうでなければ、ngにmidを代入
if A[mid] >= K:
ok = mid
else:
ng = mid
# whileの処理により、okに入っている値は、K以上を満たす最小の値。
# okが最初の値と変わっていない時、条件をみたすものがない。
if ok != N:
print(ok)
else:
print(-1)
bit全探索
sample.py
# 入力の受け取り
n = int(input())
A = list(map(int, input().split()))
q = int(input())
M = list(map(int, input().split()))
# 列挙したものを格納する配列
sum_list = []
# 1を n bit左にシフトし、上限値を1_0*(n-1) にする。
for i in range(1 << n):
l = []
# iをjbit右にシフトし、1の時だけ配列に格納するようにする
for j in range(n):
if i >> j & 1:
l.append(A[j])
# iのときの配列をもともと設定しておいた配列に格納する
# 今回は和を求めたいため、sumで格納。
sum_list.append(sum(l))
# 全通り列挙された配列をもとに、回答を表示していく。
for i in range(q):
if M[i] in sum_list:
print("yes")
else:
print("no")
順列全探索
sample.py
# ルートを使うため、今回はmathもインポートする
import itertools
import math
# 入力からNを受け取る
N = int(input())
# 座標を全て格納する
point = []
for _ in range(N):
xy = list(map(int, input().split()))
point.append(xy)
# 0 ~ N-1までの配列を作成
num = list(range(N))
# 街の番号の順列を全て列挙する
ptn = list(itertools.permutations(num))
# 合計の距離を格納する変数を定義
sum_dist = 0
# 各パターンごとに距離を求め、sum_dist に加算
for p in ptn:
for i in range(len(p)-1):
a = p[i]
b = p[i+1]
x1 = point[a][0]
y1 = point[a][1]
x2 = point[b][0]
y2 = point[b][1]
sum_dist += math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
# 加算したものを全パターンの総数で割ったもの（平均）を出力
print(sum_dist/len(ptn))
幅優先探索
sample.py
# キューをインポート
from collections import deque
# 入力を受け取る。座標の調整のため、スタート地点とゴール地点の座標を-1する。
R, C = map(int, input().split())
sy, sx = map(int, input().split())
gy, gx = map(int, input().split())
sy -= 1
sx -= 1
gy -= 1
gx -= 1
# 迷路の情報を配列Gで受け取る
G = [input() for _ in range(R)]
# キューをQに入れ、スタート地点を追加
Q = deque()
Q.append([sy, sx])
# 未訪問と始点からの距離を管理するdistを定義。スタート地点に0を代入。
dist = [[-1]*C for _ in range(R)]
dist[sy][sx] = 0
# 今回は移動する４方向を事前に用意した。
dirc = [(0, 1), (0, -1), (1, 0), (-1, 0)]
# キューの要素がなくなるまで処理を繰り返す。
while len(Q) > 0:
y, x = Q.popleft()
# 移動先で繰り返し処理
for dy, dx in dirc:
y2 = y + dy
x2 = x + dx
# 移動先が迷路の範囲内でなければ、continue
if not (0 <= y2 < R and 0 <= x2 < C):
continue
# 移動先が壁だったら、continue
if G[y2][x2] == "#":
continue
# 移動先が未訪問なら移動前の距離＋１をdistに入れて、キューに移動先の座標を追加
if dist[y2][x2] == -1:
dist[y2][x2] = dist[y][x] + 1
Q.append([y2, x2])
# ゴールの距離を出力
print(dist[gy][gx])
深さ優先探索
（幅優先探索と問題同じ）
sample.py
import sys
sys.setrecursionlimit(1000000)
## 訪問済みかどうかを管理する2次元配列
visited = []
for i in range(H):
visited.append([False]*W)
## 再帰関数dfsを定義する
def dfs(i, j):
visited[i][j] = True
# ４方向の隣マスを探索する
for i2, j2 in [[i+1, j], [i-1, j], [i, j+1], [i, j-1]]:
# もし盤面の範囲内でなければ無視する
if not (0 <= i2 < H and 0 <= j2 < W):
continue
# もし壁マスであれば無視する
if S[i][j] == "#":
continue
# もし未訪問であれば再帰的に呼び出す
if not visited[i2][j2]:
dfs(i2, j2)
## 始点から呼び出す
dfs(si, sj)
優先度付きキュー（ヒープ）
sample.py
import heapq
# リストのヒープ化
heapq.heapify(list)
# リストに追加
heapq.heappush(list, num)
# リストから取り出し
heapq.heappop(list)
ヒープから取り出せるのは、最小値のため、
大きい順に取り出したい時は、値に-1をかける。
Union-Find
sample.py
class UnionFind():
def __init__(self, n):
self.n = n
self.parents = [-1] * n
def find(self, x):
if self.parents[x] < 0:
return x
else:
self.parents[x] = self.find(self.parents[x])
return self.parents[x]
def union(self, x, y):
x = self.find(x)
y = self.find(y)
if x == y:
return
if self.parents[x] > self.parents[y]:
x, y = y, x
self.parents[x] += self.parents[y]
self.parents[y] = x
def size(self, x):
return -self.parents[self.find(x)]
def same(self, x, y):
return self.find(x) == self.find(y)
def roots(self):
return [i for i, x in enumerate(self.parents) if x < 0]
def group_count(self):
return len(self.roots())
A = [0, 1, 2, 3, 4, 5]
# 対象となる配列の要素数でインスタンスを作成
uf = UnionFind(len(A))
# グループの連結
uf.union(x, y)
# 同じグループか判定
uf.same(x, y)
細かなテクニック
リストの各要素の出現回数を得る
sample.py
from collections import Counter
l = [0, 0, 1, 1, 3, 3, 3]
cnt = Counter(l)
# Counter({3: 3, 0: 2, 1: 2})が格納され、辞書型のように使える
print(cnt[0])
# 2
小数の計算
- 小数の計算では、小数を整数に変換、計算後にもとに戻す処理を行う
- 小数→整数の変換では、掛け算より、文字列での変換が確実
例
sample.py
a = 9.79
print(int(a*100))
# 978
print(int(str(a).replace(".", "")))
# 979
マンハッタン距離の最大値
$zi = xi + yi, wi = xi - yi$
とした時、
$max( |xi-xj| + |yi-yj| )$
$= max(max(zi)-min(zi), max(wi)-min(wi))$
n>=4の包除原理
集合S1, S2, S3, ..., SN の和集合は
- 奇数個選んだとき：共通部分の要素数
- 偶数個選んだとき：共通部分の要素数×(-1)
上記の和で求められる
引用元/参考記事に関して
以下の記事やサイトを参考にさせて頂きました。