LoginSignup
11

More than 1 year has passed since last update.

競プロ用チートシート(Python)

Last updated at Posted at 2021-10-13

概要

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と設定する。それぞれ端から1ずれているのは、両端が未判定のため。
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

# 今回は移動する4方向を事前に用意した。
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

# 移動先が未訪問なら移動前の距離+1を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
  # 4方向の隣マスを探索する
  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

小数の計算

  1. 小数の計算では、小数を整数に変換、計算後にもとに戻す処理を行う
  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)
    上記の和で求められる

引用元/参考記事に関して

以下の記事やサイトを参考にさせて頂きました。

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
11