AtCoderテンプレート for Python

AtCoderテンプレート for Python

自分用のPythonAtCoderテンプレート。



入力処理

入力 x

x = int(input())

入力 x y

n,x = map(int, input().split(' '))

入力 n / a1 ... an

n = int(input())
a = list(map(int, input().split(' ')))

入力 n / s1 / s2 / ... / sn

n = int(input())
s=[]
for i in range(n):
    s.append(int(input()))

入力 n / x1 y1 / ... / xn yn

n=int(input())
pair = [tuple(map(int, input().split(' '))) for i in range(n)]


変数

配列

arr = []
print(arr)
# []

arr = [0] * 10
print(arr)
# [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

二次元配列

arr = [[]]
print(arr)
# [[]]

arr = [[0] * 4 for i in range(3)]
print(arr)
# [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]

配列の展開

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

*をつけると、配列でもタプルでもdictでも展開される。

リスト内にあるかどうか確認

import itertools

lst = [0,1,2]
print(1 in lst)
# True

print(3 in lst)
# False

lst = set([0,1,2])
print(1 in lst)
# True

lst = {0,1,2}
print(1 in lst)
# True


定数

mod = 1000000007


itertools

累積和(数値)

import itertools

arr = [0,1,3,5,7,9]
print(arr)
# [0, 1, 3, 5, 7, 9]

cumsum = list(itertools.accumulate(arr))
print(cumsum)
# [0, 1, 4, 9, 16, 25]

累積和(文字)

import itertools

s = ['ab', 'bc', 'cd']
print(s)
# ['ab', 'bc', 'cd']

scum = list(itertools.accumulate(s))
print(scum)
# ['ab', 'abbc', 'abbccd']

順列(数値)

import itertools

lst = [1, 2, 3]
print(lst)
# [1, 2, 3]

permu = list(itertools.permutations(lst))
print(permu)
# [(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]

順列(文字列)

import itertools

s = 'abc'
permus = list(itertools.permutations(s))
print(permus)
# [('a', 'b', 'c'), ('a', 'c', 'b'), ('b', 'a', 'c'), ('b', 'c', 'a'), ('c', 'a', 'b'), ('c', 'b', 'a')]

組み合わせ(数値)

import itertools

lst = [1, 2, 3, 4]
print(lst)
# [1, 2, 3, 4]

comb = list(itertools.combinations(lst, 2))
print(comb)
# [(1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)]

comb = list(itertools.combinations(lst, 3))
print(comb)
# [(1, 2, 3), (1, 2, 4), (1, 3, 4), (2, 3, 4)]

重複を許す組み合わせ(数値)

import itertools

lst = [1, 2, 3, 4]
print(lst)
# [1, 2, 3, 4]

combr = list(itertools.combinations_with_replacement(lst, 2))
print(combr)
# [(1, 1), (1, 2), (1, 3), (1, 4), (2, 2), (2, 3), (2, 4), (3, 3), (3, 4), (4, 4)]

直積

import itertools

prd = list(itertools.product([0,1], repeat=3))
print(prd)
# [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]

入力データを用いて、直積で全パターンを得る。

import itertools

lst = [(0,1), (2,3), (2,4)]
prd = list(itertools.product( *lst ))
print(prd)
# [(0, 2, 2), (0, 2, 4), (0, 3, 2), (0, 3, 4), (1, 2, 2), (1, 2, 4), (1, 3, 2), (1, 3, 4)]

for p in prd:
    p = set(p)
    print(p)
# {0, 2}
# {0, 2, 4}
# {0, 2, 3}
# {0, 3, 4}
# {1, 2}
# {1, 2, 4}
# {1, 2, 3}
# {1, 3, 4}


ビットが立っているものだけ残す

import itertools

lst = [1, 2, 3, 4]
bit = [1, 0, 0, 1]

cps = list(itertools.compress(lst, bit))
print(cps)
# [1, 4]

各値が何個連続しているか

import itertools

bi = [0,0,0,1,1,0,0,0,1,1,0,1]
print(bi)
# [0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 1]

gr = itertools.groupby(bi)
for key, group in gr:
    print(f'{key}: {list(group)}')
# 0: [0, 0, 0]
# 1: [1, 1]
# 0: [0, 0, 0]
# 1: [1, 1]
# 0: [0]
# 1: [1]

gr = itertools.groupby(bi)
for key, group in gr:
    print(f'{key}: {len(list(group))}')
# 0: 3
# 1: 2
# 0: 3
# 1: 2
# 0: 1
# 1: 1


参考:

https://docs.python.org/ja/3/library/itertools.html

https://qiita.com/SaitoTsutomu/items/ddb5076ef62745f03b56


cython

nCr:組み合わせの数(math)

import math

n = 7
r = 5
cb = math.comb(n, r)
print(cb) # -> 21


nCr:組み合わせの数(scipy) ※速い

※言語をPythonにする

from scipy.special import comb

n = 7
r = 5
cb = comb(n, r, exact=True)
print(cb) # 21

exact=Trueを指定することで戻り値が"inf"になるのを回避できる 正確な値を求めたい場合にもこの引数が必要。(その分若干遅くなる。)


上記を気にせず使いたい場合は引数を省略。

a = comb(n, r)


以下ボツ

組み合わせ

def comb(n, r):
  if r == 0 or r == n:
    return(1)
  elif r == 1:
    return(n)
  return(combination(n - 1, r - 1) + combination(n - 1, r))