【Python】AtCoderで使う基礎知識

AtCoder初心者が問題を解いていくうえで、
役立つかもしれない知識をメモとして残していく。

入力テンプレート

#---------------------------------------------------------
# N = int(input())
N,K = map(int,input().split())
# L = list(map(int,input().split()))
# L2 = list(map(int,input().split()))

# S = input()
# L = [input() for _ in range(N)]
# L = [int(input()) for _ in range(N)]
# L = [list(map(int,input().split())) for _ in range(N)]
#---------------------------------------------------------

入力

いろんな入力の受け取り。

文字の受け取り

S = input()

数値の受け取り

N = int(input())

空白数値の受け取り

A,B = map(int,input().split())

入力回数を受け取り、その回数分入力を受け取る
角かっこを使った書き方を「内包表記」という

N = int(input()) #入力回数
l = list(int(input()) for _ in range(N)) #N回の数値入力をリストとして取得
#または
l = [int(input()) for _ in range(N)] #N回の数値入力をリストとして取得

入力回数を受け取り、その回数分空白数値を受け取る

Q = int(input())
L = [0]*Q
R = [0]*Q

for i in range(Q):
    L[i], R[i] = map(int, input().split())

#または
l = list(list(map(int, input().split())) for _ in range(Q))
l = [list(map(int,(input().split()))) for _ in range(Q)]

あとで文字列を変更したい場合、文字のリストとして受け取っておく

Python #入力

T = list(input())
T[0] = "p"
print("".join(T))

python

出力

リストを空白ありで出力

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

1 2 3

リストを空白なしで出力
join

l = [1,2,3]
l = map(str,l) #joinは文字列に対してなので、文字列に変換
print("".join(l))

123

format
f文字列を文字列の前に置いておくと、変数をそのまま指定できる

name = "鈴木"
height = 170
weight = 60

print("{0}さんの身長は{1}cm、体重は{2}kgです".format(name,height,weight))
#鈴木さんの身長は170cm、体重は60kgです

print(f"{name}さんの身長は{height}cm、体重は{weight}kgです")
#鈴木さんの身長は170cm、体重は60kgです

文字列の繰り返し、文字列の乗算

s = "Python"
print(s*3)

PythonPythonPython

三項演算子

flag = 1

print("Yes" if flag else "No")
print("Yes" if not flag else "No")

Yes
No

文字列から数式を計算
eval

print(eval("1+2"))

3

ループ

2重ループで、内側ループがbreakしたかどうかの判定

for i in range(10):
    print(f'i={i}')
    for j in range(10):
        if i==3 and j==5:
            break
    else:
        print('内側ループはbreakなしで正常終了したよ')
        continue
    
    print(f'i={i} j={j}')
    print('内側ループでbreakしたよ')
    break

#i=0
#内側ループはbreakなしで正常終了したよ
#i=1
#内側ループはbreakなしで正常終了したよ
#i=2
#内側ループはbreakなしで正常終了したよ
#i=3
#i=3 j=5
#内側ループでbreakしたよ

0埋め

zfillで指定した文字数まで0埋めできる
2進数に変換して0埋めするときは、formatが便利

s = '1234'
print(s.zfill(8))
#00001234

print(format(10,'08b'))
#00001010

print(format(10,'b'))
#1010

print(bin(10))
#0b1010

各桁の和

数値を文字列に変換して、各桁の数字に分解して和を求める

x = 1357

sum(int(i) for i in str(x))
#または
sum(list(map(int, str(x))))

絶対値

abs(-2)

商と余り

x = 10; y = 3

q = x // y
mod = x % y
q, mod = divmod(x, y)

リスト

空のリストを複数作成

L= [[] for _ in range(N)]

ソート
降順にソート

l = [1,3,5,2,4]

l.sort()
l2 = sorted(l) #別変数に入れる

l.sort(reverse = True) #降順にソート
l2 = sorted(l, reverse = True) #降順にソートして別変数に入れる

2番目の要素で並び替え

l = [[1,5],[2,2],[3,4]]

l.sort(key=lambda x:x[1])
print(l)
#[[2, 2], [3, 4], [1, 5]]

スライス操作
リストを逆順に出力

l = [1,2,3,4,5]

print(l[2:]) #2番目以降
#[3, 4, 5]

print(l[:2]) #2番目まで
#[1, 2]

print(l[2:4]) #2番目以降4番目まで
#[3, 4]

print(l[-2:]) #後ろから2つ
#[4, 5]

print(l[::-1]) #リストを逆順に出力
#[5, 4, 3, 2, 1]

print(l[::2]) #1個飛ばし
#[1, 3, 5]

print(l[1::2]) #2番目以降1個飛ばし
#[2, 4]

2つのリストの各要素を足したリストを取得
operator --- 関数形式の標準演算子

from operator import *

L1 = [1,2,3]
L2 = [4,5,6]

l = list(map(add,L1,L2))
print(l)
#[5, 7, 9]

リストのすべての要素に対して、条件に合致するか判定

L1 = [2,4,6]
L2 = [2,4,7]

if all([x%2==0 for x in L1]):
    print('すべて2の倍数')
else:
    print('2の倍数でない数が存在')
#すべて2の倍数
    
if all([x%2==0 for x in L2]):
    print('すべて2の倍数')
else:
    print('2の倍数でない数が存在')
#2の倍数でない数が存在

累積和

from itertools import accumulate

l = [1, 2, 3, 4, 5, 6]
print(list(accumulate(l)))
# [1, 3, 6, 10, 15, 21]

1の位の数を取得

n1 = int(str(A)[-1])
n1 %= 10

enumerate

リストのインデックスと要素を同時に取得

l = ["国語","数学","理科","社会","英語"]

# 第2引数は、開始インデックス
for i,subject in enumerate(l, 1):
    print(str(i)+"番目は"+subject)

1番目は国語
2番目は数学
3番目は理科
4番目は社会
5番目は英語

最大公約数、最小公倍数

from math import gcd

x = 8
y = 12

gcd(x,y) #最大公約数
(x*y)//gcd(x,y) #最小公倍数

順列、組み合わせ、重複組み合わせ

from itertools import permutations, combinations, combinations_with_replacement

balls = [1, 2, 3]
print(list(permutations(balls, 3)))  #順列
#[(1, 2, 3), (1, 3, 2), (2, 1, 3), (2, 3, 1), (3, 1, 2), (3, 2, 1)]

print(list(combinations(balls, 2)))  #組み合わせ
#[(1, 2), (1, 3), (2, 3)]

print(list(combinations_with_replacement(balls,2))) #重複組み合わせ
#[(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

■重複組み合わせを使用した問題
ABC165 C - Many Requirements

階乗・順列・組み合わせ

#--階乗--------------------
from math import factorial

print(factorial(5))
# 120

print(factorial(0))
# 1

#--順列--------------------
def permutations_count(n, r):
    return math.factorial(n) // math.factorial(n - r)

print(permutations_count(4, 2))
# 12

print(permutations_count(4, 4))
# 24

#--組み合わせ---------------
from operator import mul
from functools import reduce

def combinations_count(n, r):
    r = min(r, n - r)
    numer = reduce(mul, range(n, n - r, -1), 1)
    denom = reduce(mul, range(1, r + 1), 1)
    return numer // denom

print(combinations_count(4, 2))
# 6

print(combinations_count(4, 0))
# 1

Decimal

小数の誤差が絡むような問題は、とりあえずDecimalを使ってみる。
注意点としては、Decimalに直接小数を渡す場合 文字列 を渡す！

from decimal import Decimal

print(0.3 - 0.2)
#0.09999999999999998

print(Decimal(0.3) - Decimal(0.2))
#0.09999999999999997779553950750

print(Decimal('0.3') - Decimal('0.2'))
#0.1

■使用した問題
C - Sqrt Inequality

アスキーコード ord、chr

【ord】
　アスキーコードに変換
【chr】
　アスキーコードから文字に変換

151 A - Next Alphabet

c = input()
print(chr(ord(c) + 1))

2点間の距離

dis = ((x1-x2)**2 + (y1-y2)**2)**0.5

無限大での初期化

f_inf = float('inf')
f_inf_minus = -float('inf')

重複する値を削除

l = [4,2,1,3,4,3,2]

print(set(l))

{1, 2, 3, 4}

アルファベットの出現回数を調べる

import string

val_str = "hello everyone"

d = {}

alphas = string.ascii_lowercase
for key in alphas:
    d[key] = 0

for key in val_str:
    if key in alphas:
        d[key] += 1

for i in d:
    print("{0} : {1}".format(i,d[i]))

a : 0
b : 0
c : 0
d : 0
e : 4
f : 0
g : 0
h : 1
i : 0
j : 0
k : 0
l : 2
m : 0
n : 1
o : 2
p : 0
q : 0
r : 1
s : 0
t : 0
u : 0
v : 1
w : 0
x : 0
y : 1
z : 0

出現したものだけを調べる場合

from collections import defaultdict
val_str = "hello everyone"

d = defaultdict(int)

for key in val_str:
    d[key] += 1

for i in d:
    print("{0} : {1}".format(i,d[i]))

h : 1
e : 4
l : 2
o : 2
  : 1
v : 1
r : 1
y : 1
n : 1

2分探索

146 C - Buy an Integer

A,B,Z = [int(i) for i in input().split()]

max = 10 ** 9 + 1
min = 0

while max - min > 1:
    avr = (max + min) // 2
    sum = A * avr + B * len(str(avr))
    if sum <= Z:
        min = avr
    else:
        max = avr

print(min)

スコア順にインデックスを出力

辞書(dict)を利用して出力

l = [7,6,13,10,7]

d = {}
for i,score in enumerate(l):
    d[i+1] = score

d = sorted(d.items(), key=lambda x:x[1], reverse = True)

for i in d:
    print("{0}, score:{1}".format(i[0],i[1]))

3, score:13
4, score:10
1, score:7
5, score:7
2, score:6

逆に順位を出力

l = [7,6,13,10,7]
l2 = sorted(l, reverse = True)

for i in l:
    print(l2.index(i) + 1)

3
5
1
2
3

素数判定

素数の判定には $\sqrt {n}$ まで割り切れる数があるかを確認すればよい

def is_prime(n):
    if n == 1:
        return False
    for i in range(2,int(n**0.5)+1):
        if n % i == 0:
            return False
    return True

n進数変換

引用：【Python】n進数→10進数、10進数→n進数【AtCoder】

■n進数→10進数

def base_10(num_n,n):
    num_10 = 0
    for s in str(num_n):
        num_10 *= n
        num_10 += int(s)
    return num_10

■10進数→n進数

def base_n(num_10,n):
    if num_10 == 0:
        return '0'

    str_n = ''
    while num_10:
        if num_10%n >= 10:
            return -1
        str_n += str(num_10%n)
        num_10 //= n
    return int(str_n[::-1])

■使用した問題
C - N進数

約数の一覧をリストで取得

引用：約数を高速で列挙するコード(Python)

def make_divisors(n):
    lower_divisors , upper_divisors = [], []
    i = 1
    while i*i <= n:
        if n % i == 0:
            lower_divisors.append(i)
            if i != n // i:
                upper_divisors.append(n//i)
        i += 1
    return lower_divisors + upper_divisors[::-1]

素因数分解をリストで取得

引用：Pythonで高速素因数分解〜競プロ用〜

"""nを素因数分解"""
"""2以上の整数n => [[素因数, 指数], ...]の2次元リスト"""

def factorization(n):
    arr = []
    temp = n
    for i in range(2, int(-(-n**0.5//1))+1):
        if temp%i==0:
            cnt=0
            while temp%i==0:
                cnt+=1
                temp //= i
            arr.append([i, cnt])

    if temp!=1:
        arr.append([temp, 1])

    if arr==[]:
        arr.append([n, 1])

    return arr

factorization(24)

## [[2, 3], [3, 1]] 
##  24 = 2^3 * 3^1

■使用した問題
D - Div Game

Union-Find木

引用：PythonでのUnion-Find（素集合データ構造）の実装と使い方

from collections import defaultdict
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 members(self, x):
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        group_members = defaultdict(list)
        for member in range(self.n):
            group_members[self.find(member)].append(member)
        return group_members

    def __str__(self):
        return '\n'.join(f'{r}: {m}' for r, m in self.all_group_members().items())

■使用した問題
D - Built?
B - Values

クラスカル法

※pypyで出すとメモリ制限を超えるので、pythonで提出する。

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)
 
class Kruskal:
    class Edge:
        def __init__(self, u, v, cost):
            self.u, self.v, self.cost = u, v, cost
        def __lt__(self, another):
            return self.cost < another.cost
    def __init__(self, node_size):
        self._node = node_size
        self._edge_list = []
    def add_edge(self, u, v, cost):
        self._edge_list.append(self.Edge(u, v, cost))
    def solve(self,K):
        uf = UnionFind(self._node)
        res = 0
        edge_count = 0
        self._edge_list.sort()
        for e in self._edge_list:
            if edge_count == self._node-K:
                break
            if not uf.same(e.u, e.v):
                uf.union(e.u, e.v)
                res += e.cost
                edge_count += 1
        return res

■使用した問題
finals - 本選会場 (Finals)

セグメント木

引用：【Python】セグ木、遅延セグ木【AtCoder】

#####segfunc#####
def segfunc(x,y):
    return x+y
#################
class SegTree:
    """
    init(init_val, ide_ele): 配列init_valで初期化 O(N)
    add(k, x): k番目の値にxを加算 O(logN)
    update(k, x): k番目の値をxに更新 O(logN)
    query(l, r): 区間[l, r)をsegfuncしたものを返す O(logN)
    """
    def __init__(self, init_val, segfunc, ide_ele):
        n = len(init_val)
        self.segfunc = segfunc
        self.ide_ele = ide_ele
        self.num = 1 << (n - 1).bit_length()
        self.tree = [ide_ele] * 2 * self.num
        for i in range(n):
            self.tree[self.num + i] = init_val[i]
        for i in range(self.num - 1, 0, -1):
            self.tree[i] = self.segfunc(self.tree[2 * i], self.tree[2 * i + 1])
    def add(self, k, x):
        k += self.num
        self.tree[k] += x
        while k > 1:
            self.tree[k >> 1] = self.segfunc(self.tree[k], self.tree[k ^ 1])
            k >>= 1
    def update(self, k, x):
        k += self.num
        self.tree[k] = x
        while k > 1:
            self.tree[k >> 1] = self.segfunc(self.tree[k], self.tree[k ^ 1])
            k >>= 1
    def query(self, l, r):
        res = self.ide_ele
        l += self.num
        r += self.num
        while l < r:
            if l & 1:
                res = self.segfunc(res, self.tree[l])
                l += 1
            if r & 1:
                res = self.segfunc(res, self.tree[r - 1])
            l >>= 1
            r >>= 1
        return res

#使用例
st = SegTree([0]*N,segfunc,0)

■使用した問題
C - GCD on Blackboard

めも

def numList(n):
    L=[n]
    m=n-2
    while True:
        if m<=0:
            break
        n+=m
        m-=2
        L.append(n)

    if L[0]%2==0:
        L=L+L[::-1]
    else:
        L=L+L[::-1][1:]

    return L
    

for i in range(10):
    print(i, numList(i))

# 0 [0, 0]
# 1 [1]
# 2 [2, 2]
# 3 [3, 4, 3]
# 4 [4, 6, 6, 4]
# 5 [5, 8, 9, 8, 5]
# 6 [6, 10, 12, 12, 10, 6]
# 7 [7, 12, 15, 16, 15, 12, 7]
# 8 [8, 14, 18, 20, 20, 18, 14, 8]
# 9 [9, 16, 21, 24, 25, 24, 21, 16, 9]