更新型記事です。
Python向け。

～な入力を読み込みたい

整数単独

整数単独 e.g. 12345

N = int(input())

空白区切りの整数複数(個数があらかじめ決まっているとき)

空白区切りの整数複数(個数があらかじめ決まっているとき) e.g. 12 345

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

空白区切りの整数複数

空白区切りの整数複数 e.g. 1 -2 3 -4 5 -6

A = [int(e) for e in input().split()]
# A[2] = 3
# A = [0] + [int(e) for e in input().split()] と下駄を履かせればA[1] = 1 となる

改行形式の数値リスト(N行)

改行形式の数値リスト(N行)

'''(入力例)
1
-2
3
-4
5
-6
7
'''
A = [int(input()) for _ in range(N)]    
# A[1] = -2
# A = [0] + [int(input()) for _ in range(N)] と下駄を履かせればA[1] = 1 となる

H行W列の二次元形式

H行W列の二次元形式の空白無文字

'''[(入力例)](https://atcoder.jp/contests/abc201/tasks/abc201_d)
3 3
---
+-+
+--
'''
H,W = map(int, input().split())
edge = [["$" for _ in range(W+2)]]      # "$"は問題の入力に合わせて適宜変更する。以下同様。
center = [["$"] + list(input()) + ["$"] for i in range(H)]
A = edge + center + edge
# print(A[2][1],A[2][2],A[3][2])          # + - -

H行W列の二次元形式の整数

'''(入力例) https://atcoder.jp/contests/abc201/tasks/abc201_d
3 4
1 2 3 4
5 6 7 8
9 0 1 2
'''
H,W = map(int, input().split())
edge = [["$" for _ in range(W+2)]]
center = [["$"] + [int(e) for e in input().split()] + ["$"] for i in range(H)]
A = edge + center + edge
# print(A[2][1],A[2][2],A[3][2])          # 5 6 0

N行N列の上三角行列から対角成分を除いたもの

上三角行列から対角成分を除いたもの

'''入力例
a(1,2) , a(1,3) , ... , a(1,N-1) ,a(1,N) 
a(2,3) , a(2,4) , ... , a(2,N) 
・
・
a(N-1,N)
'''
N = int(input())
H,W = N,N
edge = [["$" for _ in range(W+2)]]      # "$"は問題の入力に合わせて適宜変更する。以下同様。
center = [["$" for _ in range(i)] + [int(e) for e in input().split()] + ["$"] for i in range(2,H+1)]
a = edge + center + edge

～な形式で出力したい

リストをスペース区切り

リストをスペース区切りで

# (例)ans_list = [1,2,3,4,5] を 1 2 3 4 5 と出力したい
print(*ans_list)
# 1 2 3 4 5        print(' '.join(map(str, ans)))としても同じものが出力できる

リストを空白無しで繋げた

リストを空白無しで繋げて

# (例)ans_list = ["H","e","l","l","o"] を Hello と出力したい
print(''.join(map(str,ans_list)))

入力を取り扱いやすくしたい

アルファベットを取り扱いやすくしたい

アルファベット→数字変換

アルファベット→数字変換

print(ord(char))
# 入出力
print(ord("A"))     # 65
print(ord("B"))     # 66
# (略)
print(ord("Z"))     # 90

print(ord("a"))     # 97
# (略)
print(ord("b"))     # 98
print(ord("z"))     # 122

数字→アルファベット変換

数字→アルファベット変換

print(chr(char))
# 入出力
print(chr(65))      # A
print(chr(66))      # B
#(略)
print(chr(90))      # Z

print(chr(97))      # a
print(chr(98))      # b
#(略)
print(chr(122))     # z

alphabet = [chr(i) for i in range(97,123)]      # ['a', 'b', 'c', ... , 'x', 'y', 'z']
ALPHABET = [chr(i) for i in range(65,91)]       # ['A', 'B', 'C', ... , 'X', 'Y', 'Z']

使われている文字/数字の種類・個数をカウントしたい

collections.Counterを使う

import collections

A = [3,1,4,1,5,9,2,6,5,3,5,8]
counted_A = collections.Counter(A)

print(counted_A)                                    # Counter({5: 3, 3: 2, 1: 2, 4: 1, 9: 1, 2: 1, 6: 1, 8: 1})
print(counted_A[0],counted_A[1])                    # 0,2
counter_keys = [k for k in counted_A.keys()]        # [3, 1, 4, 5, 9, 2, 6, 8]
counter_values = [v for v in counted_A.values()]    # [2, 2, 1, 3, 1, 1, 1, 1]
counter_items = [i for i in counted_A.items()]      # [(3, 2), (1, 2), (4, 1), (5, 3), (9, 1), (2, 1), (6, 1), (8, 1)]

S = "Hello World"
counted_S = collections.Counter(S)
print(counted_S)                                    # Counter({'l': 3, 'o': 2, 'H': 1, 'e': 1, ' ': 1, 'W': 1, 'r': 1, 'd': 1})

数字の場合はこちらを使っても良い

def number_counter(n):
    number_counted = {i:0 for i in range(10)}

    if n == 0:
        number_counted[0] += 1
    else:
        while n != 0:
            tail = n % 10
            number_counted[tail] += 1
            n = n // 10

    return number_counted

print(number_counter(1234567))  # {0: 0, 1: 1, 2: 1, 3: 1, 4: 1, 5: 1, 6: 1, 7: 1, 8: 0, 9: 0}
print(number_counter(0))        # {0: 1, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0, 6: 0, 7: 0, 8: 0, 9: 0}

リストを要素とその連続する個数にまとめたい

リストの個数をまとめる(ランレングス圧縮)

import itertools

A = [int(e) for e in input().split()]   # ex. 5 7 5 7 7 8 8 9 2 1 4

RLE = [[key,len(list(group))] for key,group in itertools.groupby(A)]
#print(RLE)      # [[5, 1], [7, 1], [5, 1], [7, 2], [8, 2], [9, 1], [2, 1], [1, 1], [4, 1]]

リストの連続する数を省いて考えたい

リストから連続する数を省く

def list_compress(L):
    '''
    昇順のリストLで値が連続しているところ毎に(始,終)でまとめたものを返す
    e.g.
    [1,2,3,5,6,7,9,11,12,15,18]
    →[(1,3),(5,7),(9,9),(11,12),(15,15),(18,18)]
    '''
    compressed_L = list()

    if len(L) != 0:
        compressed_L = list()
        l = L[0]
        r = L[0]
        for i in range(1,len(L)):
            if L[i]-1 == r:
                # 今考えている区間と連続している
                r = L[i]                    # ので、区間終端を更新
            else:
                # 今考えている区間とは非連続
                compressed_L.append((l,r))
                l = L[i]
                r = L[i]
    
        compressed_L.append((l,r))
        
    return compressed_L

print(list_compress([1,2,3,5,6,7,9,11,12,15,18]))

集合・リストを処理したい

複数項目のリストを項目毎に並び替えしたい

lamda式を使う

A = [(1,3),(-1,0),(1,1),(-1,-1),(2,3),(3,1)]
# sorted(ソート対象のリスト, key=lambda x:(-x[第1優先のindex],x[第2優先のindex]))   -だと降順。なしなら昇順
print(sorted(A, key=lambda x:(-x[0],x[1]))) # [(3, 1), (2, 3), (1, 1), (1, 3), (-1, -1), (-1, 0)]
print(sorted(A, key=lambda x:(x[0],x[1])))  # [(-1, -1), (-1, 0), (1, 1), (1, 3), (2, 3), (3, 1)]
print(sorted(A, key=lambda x:(x[1],x[0])))  # [(-1, -1), (-1, 0), (1, 1), (3, 1), (1, 3), (2, 3)]

リストの左右に入れたい/取り出したい

listだと遅いのでdequeを使う

from collections import deque

A = deque([1,"two",3,"four",5])

A.append(6)         # 右端に挿入        deque([1, 'two', 3, 'four', 5, 6])
A.appendleft("#")   # 左端に挿入        deque(['#', 1, 'two', 3, 'four', 5, 6])
A.pop()             # 右端を取り出し    deque(['#', 1, 'two', 3, 'four', 5])
A.popleft()         # 左端を取り出し    deque([1, 'two', 3, 'four', 5])
print(A[0])         # 左端を表示        1
print(A[-1])        # 右端を表示        5

ランダムアクセスも行うときはごりちゃんdequeを使う。

リストの要素全体を操作したい

リストの要素数が多いとTLEするので、全体操作での合計変化を別に取っておく。

全体操作を別に分けて考慮する

N = int(input())
A = [int(e) for e in input().split()]   # [0,1,2,3,4,5]

ope_sum_for_all = 0     # 全体に行った操作での変化量合計
# 全体へ操作を行う時には ope_sum_for_all に加減算する
# A[x] + ope_sum_for_all で全体操作も考慮した値が出る

heapでも同様の考えが使える。(heapに出し入れするときに全体操作分を考慮する)

リスト全体を逆順にするのを複数回したい

回数が多いとTLEする。
反転操作回数をカウントして、偶数回なら初期と同様、奇数回なら反転状態と扱う。

e.g. 鉄則A44　https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_ar

N = int(input())
N,Q = map(int, input().split())

A = [i+1 for i in range(N)]

inverse_cnt = 0     # 反転回数のカウント

for i in range(Q):
    query = [int(e) for e in input().split()]

    if query[0] == 1:
        x,y = query[1],query[2]
        
        # 反転回数に応じて、変化させる個所を処理する
        if inverse_cnt%2 == 0:
            A[x-1] = y
        elif inverse_cnt%2 == 1:
            A[N-1-(x-1)] = y
    elif query[0] == 2:
        inverse_cnt += 1
    elif query[0] == 3:
        x = query[1]
        # 反転回数に応じて、表示個所を変える
        if inverse_cnt%2 == 0:
            print(A[x-1])
        elif inverse_cnt%2 == 1:
            print(A[N-1-(x-1)])

負の値や極大な値を含むリストを取り扱いやすくしたい

値が小さい方から0,1,2,...として取り扱えばよい (座標圧縮)

重複を排除してから変換表を作り、変換する e.g. 鉄則A15

N = int(input())
A = [int(e) for e in input().split()]           # 46 80 11 77 46

set_A = set(A)              # setにして重複を排除
list_A = list(set_A)        # listに戻す
list_A.sort(reverse=False)  # 小さい順に並び替え

# 小さい方から1,2,～を割り当てるように、変換テーブルを作成
conv_dict = dict()      # key:オリジナルの値 , value:座標圧縮後の値
inv_conv_dict = dict()  # key:座標圧縮後の値 , value:オリジナルの値

for i in range(len(list_A)):
    conv_dict[list_A[i]] = i
    inv_conv_dict[i] = list_A[i]

# 変換テーブルに従って変換を実施
compressed_A = [conv_dict[A[i]] for i in range(N)]

print(*compressed_A)                                       # 1 3 0 2 1

リストの要素の全取り組み合わせに対して総和や最大値を求める問題で、i < jな制限がある

二重シグマが出てくるような問題。

後ろから処理すれば自動的に満たす。 cf.ABC351-F
処理が順序を前後しても同じ(e.g. 加算,掛け算)ならばsortを掛けてから処理する。

リストの転倒数を求めたい

転倒数･･･ある数列を、隣り合う数字を並び替える操作のみで、昇順に並び替えるのに必要な操作回数

座標圧縮してから、前の方からBITで処理

def get_inv_number(A):    
    ''' 1次元リストAの転倒数を求める '''
    
    # 1次元リストAを座標圧縮する
    N = len(A)
    single_listed_A = list(set(A))
    single_listed_A.sort(reverse=False)
    conv_dict = dict()      # key:オリジナルの値 , value:座標圧縮後の値
    inv_conv_dict = dict()  # key:座標圧縮後の値 , value:オリジナルの値

    for i in range(len(single_listed_A)):
        conv_dict[single_listed_A[i]] = i
        inv_conv_dict[i] = single_listed_A[i]

    converted_A = [conv_dict[A[i]] for i in range(N)]   # 座標圧縮した結果
    
    # BITを求めて転倒数を求める
    fenwick_tree = FenwickTree(N+1)   # 0～Nまでの要素を用いることができるBIT木を確保
    inv_number = 0

    for i in range(N):
        inv_number += fenwick_tree.sum(converted_A[i], N+1)   # BIT木の[x～y)の和を算出する。(y番目の要素は含まないことに注意!)
        # ↑ 自分よりリスト左側にあるのに、値が大きいものの個数を求めている
        fenwick_tree.add(converted_A[i], 1)                   # BIT木のx番目の要素にyを加算する 

    return inv_number

集合の中から最小(最大)値を順次取り出し・参照したい

値の追加がない場合

sortしてpopする

A = [3,-1,4,-1,5,-9,26]

# 最小値の場合
A.sort(reverse=True)
A_min = A[-1]       # 参照
A_min = A.pop(-1)   # 取り出し

# 最大値の場合
A.sort(reverse=False)
A_max = A[-1]       # 参照
A_max = A.pop(-1)   # 取り出し

# 最小値/最大値 両方の取り出しが必要な場合 → dequeを使う
from collections import deque
A.sort(reverse=False)
A = deque(A)
A_min = A[0]        # 最小値参照
A_max = A[-1]       # 最大値参照
A_min = A.popleft() # 最小値取り出し
A_max = A.pop()     # 最大値取り出し

値の追加がある場合

毎回sortしていると遅いのでheapqを使う

最小値取り出しの場合

import heapq
A = [3,-1,4,-1,5,-9,26]

heapq.heapify(A)                # Aを優先度付きキューに変換する
heapq.heappush(A,-53)           # Aに値を挿入する ※型に注意。strで読んだときはintに直して入れること
A_min = heapq.heappop(A)        # Aの最小値を取り出す   A_min = -53
print(A[0])                     # Aの最小値を参照する   -9

最大値取り出しの場合

import heapq
A = [3,-1,4,-1,5,-9,26]

# heapqは最小値優先なので-1倍することで最大値優先にする
B = [-A[i] for i in range(len(A))]              # 最小値優先であるheapqを最大値取り出しで使うために-1倍しておく

heapq.heapify(B)                                # 優先度付きキューに変換する

heapq.heappush(B,-1 * -53)                      # 値(-53)を挿入する。入れるときには-1倍して入れる。　※型に注意。strで読んだときはintに直して入れること
A_max = -1 * heapq.heappop(B)                   # 最大値を"取り出す" 出すときにも-1倍して正負を戻す。   A_max = 26
print(-1 * B[0])                                # 最大値を"参照する" 出すときにも-1倍して正負を戻す。   5

値の変化・追加がある場合

問題例 : ABC267-E「Erasing Vertices 2」
→ プライオリティーキューを用いた差分更新を行う。
　値と取り出し未/済をリストで管理しつつ、(値,index)をheapqに突っ込むことで、リスト中の最小値を高速に取り出せる。

プライオリティーキューを用いた差分更新

import heapq

class VariablePriorityQueue:
    def __init__(self, v=None , reverse=False):
        self._n = len(v)
        self._v = v
        self._rev = reverse

        if v is not None:
            if not self._rev:
                self._d = [(v[i],i) for i in range(self._n)]
            else:
                self._d = [(-v[i],i) for i in range(self._n)]
        else:
            self._d = []
        
        heapq.heapify(self._d)
        
    def add(self, i, x):
        # 加算クエリ: i番目の数にxを加算
        if not self._rev:
            self._v[i] += x
            heapq.heappush(self._d,(self._v[i],i))
        else:
            self._v[i] -= x
            heapq.heappush(self._d,(-self._v[i],i))

    def append(self, x):
        # 追加クエリ: 数xを追加
        if not self._rev:
            self._v.append(x)
            heapq.heappush(self._d,(x,self._n))
        else:
            self._v.append(-x)
            heapq.heappush(self._d,(-x,self._n))

        self._n += 1
    
    def get(self):
        # 取得クエリ: 最小値or最大値を取得
        while len(self._d) != 0:
            x,i = heapq.heappop(self._d)
            if self._v[i] == x:
                if not self._rev:
                    return x
                else:
                    return -x
    def ref(self):
        # 参照クエリ: 最小値or最大値を確認
        while len(self._d) != 0:
            x,i = self._d[0]
            if self._v[i] == x:
                if not self._rev:
                    return x
                else:
                    return -x
            else:
                heapq.heappop(self._d)

# 入力読み込み
A = [int(e) for e in input().split()]     # [3 1 4 1 5 9]

VPQ = VariablePriorityQueue(A)
print(VPQ.ref())    # 1               [3 1 4 1 5 9]
VPQ.add(1,5)        # [3 1 4 1 5 9] → [3 6 4 1 5 9]
VPQ.add(3,1)        # [3 6 4 1 5 9] → [3 6 4 2 5 9]
print(VPQ.ref())    # 2
VPQ.append(-10)     # [3 6 4 1 5 9] → [3 6 4 2 5 9 -10]
print(VPQ.get())    # -10             [3 6 4 2 5 9 ___]
print(VPQ.get())    # 2               [3 6 4 _ 5 9 ___]

値の追加・削除・変更がある場合

問題例 : 鉄則A38「Black Company 1」
→ プライオリティーキューを用いつつ、削除/変更された値をdictやlistで管理しておく。

値の追加・削除がある場合の最小値取り出し

import heapq
from collections import defaultdict

D,N = map(int, input().split())

limit_on = [list() for i in range(D+2)]
limit_off = [list() for i in range(D+2)]

for i in range(N):
    L,R,H = map(int, input().split())
    limit_on[L].append(H)
    limit_off[R+1].append(H)

limit_hq = [24]                     # 最小(かもしれない)値が格納されるheap que。問題文より制約分を入れておく
heapq.heapify(limit_hq)             # 最小値取り出しの為にheap化

invalid_limit = defaultdict(int)    # 削除された値とその個数を管理

ans = 0
for i in range(1,D+1):
    for H in limit_on[i]:
        heapq.heappush(limit_hq,H)
    for H in limit_off[i]:
        invalid_limit[H] += 1

    while len(limit_hq) != 0:
        limit = limit_hq[0]

        if invalid_limit[limit] > 0:
            # 当該値は削除されているので最小値としては不適切
            heapq.heappop(limit_hq)     # 最小値候補から除外
            invalid_limit[limit] -= 1   # 削除した分カウントを減らす
        else:
            break

    ans += limit

print(ans)

変化のないリストの中から～未満/以下/より大きい/以上な数に一番近い数や個数を調べたい

bisectを使う

import bisect
# index = bisect.bisect_right(A_list,insert_number)   #場所を求める
# bisect.insort_left(A_list,insert_number)            #挿入する

INF = pow(10,6)
A = [-10,3,3,4,7,9,11]          # 対象のリストは .sort(reverse) すること!!
X = int(input())
Y = int(input())

# ◆ sort済みのAリスト中でX以下の最大値を求める
index = bisect.bisect_right(A,X)
if index == 0:
    v = -INF
else:
    v = A[index-1]

print(v,index,len(A)-index)     # X以下の最大値,X以下の個数,Xより大きい個数

# ◆ sort済みのAリスト中でX未満の最大値を求める
index = bisect.bisect_left(A,X)
if index == 0:
    v = -INF
else:
    v = A[index-1]

print(v,index,len(A)-index)     # X未満の最大値,X未満の個数,X以上の個数

# ◆ sort済みのAリスト中でXより大きい数のうち最小値を求める
index = bisect.bisect_right(A,X)
if index == len(A):
    v = INF
else:
    v = A[index]

print(v,index,len(A)-index)      # Xより大きい数のうち最小値,X以下の個数,Xより大きいの個数

# ◆ sort済みのAリスト中でX以上の最小値を求める
index = bisect.bisect_left(A,X)
if index == len(A):
    v = INF
else:
    v = A[index]

print(v,index,len(A)-index)     # X以上の最小値,X未満の個数,X以上の個数

# ◆ sort済みのAリスト中でX以上Y以下を抽出/個数を求める
if X > Y:
    print("input error")        # X以上Y以下なので入力制約に引っかかっている
else:
    index_left = bisect.bisect_left(A,X)

    if index_left == len(A):
        print(list(),0)         # X以上の値が無いので終了
    else:
        index_right = bisect.bisect_right(A,Y)

        print(A[index_left:index_right],index_right - index_left)

# ◆ sort済みのAリスト中でXと絶対値が最も小さくなる数を求める
close_value = -INF

index = bisect.bisect_left(A,X)
if index != 0:
    v = A[index-1]              # sort済みのAリスト中でX未満の最大値を求める
    close_value = v

index = bisect.bisect_left(A,X)
if index != len(A):
    v = A[index]                # sort済みのAリスト中でX以上の最小値を求める
    if abs(close_value - X) > abs(v - X):
        close_value = v

print(close_value)

変化のないリストの中から～未満/以下/より大きい/以上な/指定範囲にある個数を調べたい

何回も調べる必要があるならこちらで。

累積和を使う

import collections

INF = pow(10,6)
A = [-10,3,3,4,7,9,11]

# ●Aにある値以下が何個あるかを求めて置く。(累積和)
cum_sum = [0] * (pow(10,6)+2)
counted_A = collections.Counter(A)
for i in range(1,pow(10,6)+2):
    cum_sum[i] = cum_sum[i-1] + counted_A[i]

# ◆ Aリスト中でX以下の個数を求める
print(cum_sum[X])

# ◆ Aリスト中でX未満の個数を求める
print(cum_sum[X-1])

# ◆ Aリスト中でX以上Y以下の個数を求める
print(cum_sum[X] - cum_sum[Y-1])

変化のあるリストの中から～未満/以下/より大きい/以上な数に一番近い数や個数調べたい

平方分割を利用したSorted (Multi) Setを使う。

Sorted Multi Set

# 本体   →  https://github.com/tatyam-prime/SortedSet/blob/main/SortedMultiset.py
# 使用法 →  https://github.com/tatyam-prime/SortedSet
# 解説   →  https://qiita.com/tatyam/items/492c70ac4c955c055602
import math
from bisect import bisect_left, bisect_right
from typing import Generic, Iterable, Iterator, List, Tuple, TypeVar, Optional
T = TypeVar('T')

class SortedMultiset(Generic[T]):
    BUCKET_RATIO = 50
    REBUILD_RATIO = 170

    def _build(self, a: Optional[List[T]] = None) -> None:
        "Evenly divide `a` into buckets."
        if a is None: a = list(self)
        size = len(a)
        bucket_size = int(math.ceil(math.sqrt(size / self.BUCKET_RATIO)))
        self.a = [a[size * i // bucket_size : size * (i + 1) // bucket_size] for i in range(bucket_size)]
    
    def __init__(self, a: Iterable[T] = []) -> None:
        "Make a new SortedMultiset from iterable. / O(N) if sorted / O(N log N)"
        a = list(a)
        self.size = len(a)
        if not all(a[i] <= a[i + 1] for i in range(len(a) - 1)):
            a = sorted(a)
        self._build(a)

    def __iter__(self) -> Iterator[T]:
        for i in self.a:
            for j in i: yield j

    def __reversed__(self) -> Iterator[T]:
        for i in reversed(self.a):
            for j in reversed(i): yield j
    
    def __eq__(self, other) -> bool:
        return list(self) == list(other)
    
    def __len__(self) -> int:
        return self.size
    
    def __repr__(self) -> str:
        return "SortedMultiset" + str(self.a)
    
    def __str__(self) -> str:
        s = str(list(self))
        return "{" + s[1 : len(s) - 1] + "}"

    def _position(self, x: T) -> Tuple[List[T], int]:
        "Find the bucket and position which x should be inserted. self must not be empty."
        for a in self.a:
            if x <= a[-1]: break
        return (a, bisect_left(a, x))

    def __contains__(self, x: T) -> bool:
        if self.size == 0: return False
        a, i = self._position(x)
        return i != len(a) and a[i] == x

    def count(self, x: T) -> int:
        "Count the number of x."
        return self.index_right(x) - self.index(x)

    def add(self, x: T) -> None:
        "Add an element. / O(√N)"
        if self.size == 0:
            self.a = [[x]]
            self.size = 1
            return
        a, i = self._position(x)
        a.insert(i, x)
        self.size += 1
        if len(a) > len(self.a) * self.REBUILD_RATIO:
            self._build()
    
    def _pop(self, a: List[T], i: int) -> T:
        ans = a.pop(i)
        self.size -= 1
        if not a: self._build()
        return ans

    def discard(self, x: T) -> bool:
        "Remove an element and return True if removed. / O(√N)"
        if self.size == 0: return False
        a, i = self._position(x)
        if i == len(a) or a[i] != x: return False
        self._pop(a, i)
        return True

    def lt(self, x: T) -> Optional[T]:
        "Find the largest element < x, or None if it doesn't exist."
        for a in reversed(self.a):
            if a[0] < x:
                return a[bisect_left(a, x) - 1]

    def le(self, x: T) -> Optional[T]:
        "Find the largest element <= x, or None if it doesn't exist."
        for a in reversed(self.a):
            if a[0] <= x:
                return a[bisect_right(a, x) - 1]

    def gt(self, x: T) -> Optional[T]:
        "Find the smallest element > x, or None if it doesn't exist."
        for a in self.a:
            if a[-1] > x:
                return a[bisect_right(a, x)]

    def ge(self, x: T) -> Optional[T]:
        "Find the smallest element >= x, or None if it doesn't exist."
        for a in self.a:
            if a[-1] >= x:
                return a[bisect_left(a, x)]
    
    def __getitem__(self, i: int) -> T:
        "Return the i-th element.  (i番目の要素の値を返す)"
        if i < 0:
            for a in reversed(self.a):
                i += len(a)
                if i >= 0: return a[i]
        else:
            for a in self.a:
                if i < len(a): return a[i]
                i -= len(a)
        raise IndexError
    
    def pop(self, i: int = -1) -> T:
        "Pop and return the i-th element. (i番目の要素を削除して値を返す)"
        if i < 0:
            for a in reversed(self.a):
                i += len(a)
                if i >= 0: return self._pop(a, i)
        else:
            for a in self.a:
                if i < len(a): return self._pop(a, i)
                i -= len(a)
        raise IndexError

    def index(self, x: T) -> int:
        "Count the number of elements < x. (x未満の要素個の個数を調べる)"
        ans = 0
        for a in self.a:
            if a[-1] >= x:
                return ans + bisect_left(a, x)
            ans += len(a)
        return ans

    def index_right(self, x: T) -> int:
        "Count the number of elements <= x. (x以下の要素の個数を調べる)"
        ans = 0
        for a in self.a:
            if a[-1] > x:
                return ans + bisect_right(a, x)
            ans += len(a)
        return ans

# 使用例
A = [int(e) for e in input().split()]   # [3 -1 4 -1 5 -9 2 -6 5 -3 5 -8]

s = SortedMultiset(A)                   # iterable から SortedMultiSet を作る
print(s.a)                              # s.a で SortedMultiSet の中身を見れる。listのlistなので合計などを取るときには[0]を付けること。
#>> [[-9, -8, -6, -3, -1, -1, 2, 3, 4, 5, 5, 5]]
s.add(0)                                # 値を追加 [[-9, -8, -6, -3, -1, -1, 0, 2, 3, 4, 5, 5, 5]]
s.discard(-1)                           # 値を削除 [[-9, -8, -6, -3, -1, 0, 2, 3, 4, 5, 5, 5]]

print(s.count(5))                       # 値の個数を調べる                      3

print(s.lt(2))                          # 指定値より小さい最大要素を調べる      0
print(s.le(2))                          # 指定値以下の最大要素を調べる          2
print(s.gt(2))                          # 指定値より大きい最小要素を調べる      3
print(s.ge(2))                          # 指定値以上の最小要素を調べる          2

print(s.index_right(2))                 # 指定値以下の要素の個数を調べる        7
print(s.index(2))                       # 指定値未満の要素の個数を調べる        6
print(len(s) - s.index(2))              # 指定値以上の要素の個数を調べる        6
print(len(s) - s.index_right(2))        # 指定値より大きいの要素の個数を調べる  5

途中変化なしなリストの区間の合計を求めたい

区間和問題
e.g. 鉄則B06…回答

累積和を使う

import itertools
A = [1, 2, 3, 4, 5, 6]   # 以下の処理は Aが A[1] A[2] ... と与えられたときに基づく

cumsum_A = [0] + list(itertools.accumulate(A))      # 累積和の作成(1-index)
print(cumsum_A)     # [0, 1, 3, 6, 10, 15, 21]

L,R = 3,5
print(cumsum_A[R] - cumsum_A[L-1],R-L+1)            # L～Rの区間和と区間内個数         12 , 3
print(cumsum_A[-1] - cumsum_A[L-1],len(A)-L+1)      # L～最後までの区間和と区間内個数  18 , 4
print(cumsum_A[R],R)                                # 最初～Rまでの区間和と区間内個数  15 , 5

幅が一定な区間を頭から順番に求めるなら、dequeを用いて、新しく範囲内になった数を合計に加え、範囲外になった数を合計から引いても良い。

途中変化なしな区域の範囲の合計を求めたい

e.g. 鉄則A08…回答

二次元累積和を使う

H,W = map(int, input().split())
edge = [[0 for _ in range(W+2)]]
center = [[0] + [int(e) for e in input().split()] + [0] for i in range(H)]
X = edge + center + edge

# 二次元累積和を求める
cumsum_2d = [[0 for _ in range(W+2)] for _ in range(H+2)]   # 二次元累積和
for i in range(1,H+1):
    cumsum_2d[i][0] = 0
    for j in range(1,W+1):
        cumsum_2d[i][j] = cumsum_2d[i-1][j] + cumsum_2d[i][j-1] - cumsum_2d[i-1][j-1] + X[i][j]

# クエリに回答する
Q = int(input())
for i in range(Q):
    A,B,C,D = map(int, input().split())
    print(cumsum_2d[C][D] - cumsum_2d[A-1][D] - cumsum_2d[C][B-1] + cumsum_2d[A-1][B-1])

途中変化なしな領域の範囲の合計を求めたい

e.g. ABC366_D

3次元累積和を使う

N = int(input())
A = [[[int(e) for e in input().split()] for _ in range(N)] for _ in range(N)]

cumsum_3d = [[[0 for _ in range(N+1)] for _ in range(N+1)] for _ in range(N+1)]

for i in range(1, N+1):
    for j in range(1, N+1):
        for k in range(1, N+1):
            cumsum_3d[i][j][k] = (
                cumsum_3d[i-1][j][k] +
                cumsum_3d[i][j-1][k] +
                cumsum_3d[i][j][k-1] -
                cumsum_3d[i-1][j-1][k] -
                cumsum_3d[i-1][j][k-1] -
                cumsum_3d[i][j-1][k-1] +
                cumsum_3d[i-1][j-1][k-1] +
                A[i-1][j-1][k-1]
            )

# クエリに回答する
Q = int(input())
for _ in range(Q):
    Lx, Rx, Ly, Ry, Lz, Rz = map(int, input().split())
    result = (
        cumsum_3d[Rx][Ry][Rz] -
        cumsum_3d[Lx-1][Ry][Rz] -
        cumsum_3d[Rx][Ly-1][Rz] -
        cumsum_3d[Rx][Ry][Lz-1] +
        cumsum_3d[Lx-1][Ly-1][Rz] +
        cumsum_3d[Lx-1][Ry][Lz-1] +
        cumsum_3d[Rx][Ly-1][Lz-1] -
        cumsum_3d[Lx-1][Ly-1][Lz-1]
    )
    print(result)

途中変化なしな区間の要素が一致するかを求めたい

e.g. ABC367_F

ハッシュして累積和を使う

import random

rand_max = pow(2,64)

# 入力受取
N,Q = map(int, input().split())
A = [int(e) for e in input().split()]
B = [int(e) for e in input().split()]
set_A = set(A)
set_B = set(B)

# A,Bに現れる数字に適当に重ならないような数字(hash)を割り振る
hash_encode = dict()
hash_decode = dict()

for i in set_A:
    while 1:
        rand_num = random.randrange(rand_max)

        if rand_num not in hash_decode:
            hash_encode[i] = rand_num
            hash_decode[rand_num] = i
            break

for i in set_B:
    if i not in set_A:
        while 1:
            rand_num = random.randrange(rand_max)

            if rand_num not in hash_decode:
                hash_encode[i] = rand_num
                hash_decode[rand_num] = i
                break


# hashした結果を用いた累積和を求める
cumsum_A = [0]
cumsum_B = [0]
for i in range(N):
    cumsum_A.append(cumsum_A[-1] + hash_encode[A[i]])
    cumsum_B.append(cumsum_B[-1] + hash_encode[B[i]])


# クエリに回答する
for i in range(Q):
    l,r,L,R = map(int, input().split())

    if cumsum_A[r] - cumsum_A[l-1] == cumsum_B[R] - cumsum_B[L-1]:
        print("Yes")
    else:
        print("No")

区間に対する増減量を複数与えられたときに、最終時点での各点での合計量を求めたい

imos法(いもす法)を使う e.g.第14回日本情報オリンピック本選（オンライン）-A

# 入力読み込み
N,M = map(int, input().split())
P = [int(e) for e in input().split()]

ABC = [0]
for i in range(N-1):
    A,B,C = map(int, input().split())
    ABC.append((A,B,C))


# 各鉄道の利用回数をimos法を用いてまとめる
# まずは区間に対する増減をまとめる
variations = [0 for i in range(N+2)]    # imos法で用いるための増減の記録先。

for i in range(1,M):
    Pa,Pb = P[i-1],P[i]
    if Pa > Pb:
        Pa,Pb = Pb,Pa
    variations[Pa] += 1     # 利用区間の始まり   で+
    variations[Pb] -= 1     # 利用区間の終わり+1 で- (今回はPa,Pbが頂点、variationsは辺で考えているので、Pb+1としなくてもOK)

# まとめた結果を利用して、imos法で最終時点での値を求める
use_cnt = [variations[0]]

for i in range(1,N):
    use_cnt.append(use_cnt[-1] + variations[i])


# 各鉄道の利用回数に応じて、どちらが安いかを判断して、足していく
ans = 0
for i in range(1,N):
    A,B,C = ABC[i]
    
    if A * use_cnt[i] > C + B * use_cnt[i]:
        ans += C + B * use_cnt[i]
    else:
        ans += A * use_cnt[i]

print(ans)

区間に対する増減量を複数与えられたときに、途中途中での各点での量を求めたい

imos法(いもす法)とセグメント木を組み合わせる

# Problem ･･･ https://atcoder.jp/contests/abc340/tasks/abc340_e

class SegTree:
    # https://qiita.com/R_olldIce/items/32cbf5bc3ffb2f84a898
    def __init__(self, op, e, n, v=None):
        self._n = n
        self._op = op
        self._e = e
        self._log = (n - 1).bit_length()
        self._size = 1 << self._log
        self._d = [self._e()] * (self._size << 1)
        if v is not None:
            for i in range(self._n):
                self._d[self._size + i] = v[i]
            for i in range(self._size - 1, 0, -1):
                self._d[i] = self._op(self._d[i << 1], self._d[i << 1 | 1])

    def set(self, p, x):
        # 更新クエリ: a[p] を x に更新
        p += self._size
        self._d[p] = x
        while p:
            self._d[p >> 1] = self._op(self._d[p], self._d[p ^ 1])
            p >>= 1

    def get(self, p):
        # p が与えられたとき、a[p] を取得する
        return self._d[p + self._size]

    def prod(self, l, r):
        # 取得クエリ: 区間 [l,r) の総積を取得 (※ここでの積は"掛け算"でなく"演算結果"くらいの意味)
        # 【注意!】: ")" なので区間の終わりの数字ちょうどは含まれない! 必要に応じて+1すること!
        sml, smr = self._e(), self._e()
        l += self._size
        r += self._size
        while l < r:
            if l & 1:
                sml = self._op(sml, self._d[l])
                l += 1
            if r & 1:
                r -= 1
                smr = self._op(self._d[r], smr)
            l >>= 1
            r >>= 1
        return self._op(sml, smr)

    def all_prod(self):
        # 全要素の総積を取得 (※ここでの積は"掛け算"でなく"演算結果"くらいの意味)
        return self._d[1]

def op(x, y):
    # 演算
    return x + y

def e():
    # 単位元。任意の要素に対して演算をしても影響を与えないような値。
    # opが加算なら0,掛け算なら1,minならINF,maxなら-INF,xorなら0,andなら1,orなら0
    return 0
    
# 入力読み込み
N,M = map(int, input().split())
A = [int(e) for e in input().split()]
B = [int(e) for e in input().split()]

# セグ木の確保と初期化
seg = SegTree(op, e, N+1)    # 左から演算,単位元,初期値リストの長さ,初期値リスト。imos法のようにボールの増減を入れる
for i in range(N):
    seg.set(i,seg.get(i) + A[i])        # segにはimos法で記入するので、箱のボール数分だけ+
    seg.set(i+1,seg.get(i+1) - A[i])    # segにはimos法で記入するので、その隣には箱のボール数分だけ-

# 操作
for i in range(M):
    # 箱のボールを0にする
    box_num = B[i]
    ball_cnt = seg.prod(0,box_num+1)                    # segにはimos法で記入しているので、これでbox_numのボール個数が出る
    seg.set(box_num,seg.get(box_num) - ball_cnt)
    seg.set(box_num+1,seg.get(box_num+1) + ball_cnt)

    # まずbox_num+1 ~ N-1の箱までボールを1個ずつ入れることを試みる
    if box_num != N-1:
        # 取り出し元の箱番号がN-1でなければ、できるか考えてみる
        if (N - 1) - (box_num + 1) + 1 <= ball_cnt:
            # OKなら入れる
            seg.set(box_num + 1,seg.get(box_num + 1) + 1)
            seg.set(N,seg.get(N) - 1)
            ball_cnt -= (N - 1) - (box_num + 1) + 1
        else:
            # NGなら入れられるところまで入れる
            seg.set(box_num + 1,seg.get(box_num + 1) + 1)
            seg.set(box_num + ball_cnt + 1,seg.get(box_num + ball_cnt + 1) - 1)
            ball_cnt = 0
    
    # 次に0 ~ N-1の箱まで全てにボールを同じ個数個入れる
    j = ball_cnt // N
    seg.set(0,seg.get(0) + j)
    seg.set(N,seg.get(N) - j)
    ball_cnt = ball_cnt % N

    # 最後に0～入れられるところまでボールを1個ずつ入れる    
    if ball_cnt != 0:
        seg.set(0,seg.get(0) + 1)
        seg.set(ball_cnt,seg.get(ball_cnt) - 1)

# 操作後の個数を求める
ans = [seg.prod(0,i+1) for i in range(N)]
print(*ans)

領域の変化量を与えられたときに、最終的な量を求めたい

e.g. 鉄則A09…回答

2次元imos法(いもす法)を使う

# cf: https://imoz.jp/algorithms/imos_method.html

H,W,N = map(int, input().split())

weight = [[0 for _ in range(W+2)] for _ in range(H+2)]
ans = [[0 for _ in range(W+2)] for _ in range(H+2)]

# 重みの加算をする
for i in range(N):
    A,B,C,D = map(int, input().split())
    weight[A][B] += 1
    weight[A][D+1] -= 1
    weight[C+1][B] -= 1
    weight[C+1][D+1] += 1

# 横方向に累積する
for i in range(1,H+1):
    imos = 0
    for j in range(1,W+1):
        imos += weight[i][j]
        weight[i][j] = imos

# 縦方向に累積する
for j in range(1,W+1):
    imos = 0
    for i in range(1,H+1):
        imos += weight[i][j]
        ans[i][j] = imos

# 回答を出力
for i in range(1,H+1):
    print(*ans[i][1:W+1])

途中変化ありなリストの任意区間の合計値を求めたい

BIT(Fenwick Tree)を使う

import typing

class FenwickTree:
    # https://github.com/not522/ac-library-python/blob/master/atcoder/fenwicktree.py
    '''Reference: https://en.wikipedia.org/wiki/Fenwick_tree'''

    def __init__(self, n: int = 0) -> None:
        self._n = n
        self.data = [0] * n

    def add(self, p: int, x: typing.Any) -> None:
        assert 0 <= p < self._n

        p += 1
        while p <= self._n:
            self.data[p - 1] += x
            p += p & -p

    def sum(self, left: int, right: int) -> typing.Any:
        assert 0 <= left <= right <= self._n

        return self._sum(right) - self._sum(left)

    def _sum(self, r: int) -> typing.Any:
        s = 0
        while r > 0:
            s += self.data[r - 1]
            r -= r & -r

        return s
    
N = len(p)
fenwick_tree = FenwickTree(N+1)     # 0～Nまでの要素を用いることができるBIT木を確保
fenwick_tree.sum(x,y)               # BIT木の[x～y)の和を算出する。(y番目の要素は含まないことに注意!)
fenwick_tree.add(x,y)               # BIT木のx番目の要素にyを加算する (加算であり書き換えでないことに注意!)
# 書き換える場合は値を調べて、その分をマイナスしてから加算する
v = fenwick_tree.sum(x,x+1)
fenwick_tree.add(x,-v+y)            # BIT木のx番目の要素をyに書き換える

途中変化ありなリストの任意区間の合計/最高値/最低値 etcを求めたい

セグメント木を使う

# https://qiita.com/R_olldIce/items/32cbf5bc3ffb2f84a898
class SegTree:
    def __init__(self, op, e, n, v=None):
        self._n = n
        self._op = op
        self._e = e
        self._log = (n - 1).bit_length()
        self._size = 1 << self._log
        self._d = [self._e()] * (self._size << 1)
        if v is not None:
            for i in range(self._n):
                self._d[self._size + i] = v[i]
            for i in range(self._size - 1, 0, -1):
                self._d[i] = self._op(self._d[i << 1], self._d[i << 1 | 1])

    def set(self, p, x):
        # 更新クエリ: a[p] を x に更新
        p += self._size
        self._d[p] = x
        while p:
            self._d[p >> 1] = self._op(self._d[p], self._d[p ^ 1])
            p >>= 1

    def get(self, p):
        # p が与えられたとき、a[p] を取得する
        return self._d[p + self._size]

    def prod(self, l, r):
        # 取得クエリ: 区間 [l,r) の総積を取得 (※ここでの積は"掛け算"でなく"演算結果"くらいの意味)
        # 【注意!】: ")" なので区間の終わりの数字ちょうどは含まれない! 必要に応じて+1すること!
        sml, smr = self._e(), self._e()
        l += self._size
        r += self._size
        while l < r:
            if l & 1:
                sml = self._op(sml, self._d[l])
                l += 1
            if r & 1:
                r -= 1
                smr = self._op(self._d[r], smr)
            l >>= 1
            r >>= 1
        return self._op(sml, smr)

    def all_prod(self):
        # 全要素の総積を取得 (※ここでの積は"掛け算"でなく"演算結果"くらいの意味)
        return self._d[1]

def op(x, y):
    # 演算
    return x + y

def e():
    # 単位元。任意の要素に対して演算をしても影響を与えないような値。
    # opが加算なら0,掛け算なら1,minならINF,maxなら-INF,xorなら0,andなら1,orなら0
    return 0

# 使用例
N = int(input())        # 5
init_list = [i for i in range(N)]                   # 初期値リスト。問題に合わせて変更。今回は[0,1,2,...,N-1]とした。
# init_list = [0] *(N)                              # 0埋めなら、このように書くと速い
seg = SegTree(op, e, len(init_list) , init_list)    # 左から演算,単位元,初期値リストの長さ,初期値リスト

print(seg.get(3))       # 葉3の値を取得。 = 3
print(seg.prod(1,4))    # 葉1~葉3の演算(和)結果を取得。= 1 + 2 + 3 = 6
# 【注意！】↑ 開区間なので終わりの方には+1をする!!

seg.set(3 , 6)          # 葉3の値を6に更新
print(seg.prod(1,4))    # 葉1~葉3の演算(和)結果を取得。1 + 2 + 6 = 9

各要素で個数を取っておき、和を取るようにすれば、～以下の個数や～以上の個数を求めるのに利用可能
各要素の最大値が大きい・種類がそれほどでもない場合は座標圧縮と組み合わせる。 cf.ABC351-F

途中変化なしなリストの一定区間の合計/最高値/最低値 etcを求めたい

スライド最小値を使う

import heapq
from collections import deque
from collections import defaultdict

N = int(input())                        # Aの長さ
A = [int(e) for e in input().split()]
K = int(input())                        # 取りたい区間長さ
deq = deque([])                         # 一定区間に入っている値が入るdeque
heap_que = list()           
heapq.heapify(heap_que)                 # deqにある可能性がある値が入るheap
valid_value_cnt = defaultdict(int)      # deqの値とその個数を管理するdict

for i in range(N):
    deq.append(A[i])
    valid_value_cnt[A[i]] += 1
    heapq.heappush(heap_que,A[i]) 
    
    # deqが区間長の制約を超えていたら、最古値を削除する
    if len(deq) > K:
        v = deq.popleft()
        valid_value_cnt[v] -= 1

    # deqの最小値をheap_queを用いて取り出す
    while len(heap_que) != 0:
        v = heap_que[0]
        if valid_value_cnt[v] == 0:
            heapq.heappop(heap_que)
        else:
            break

cf. スライド最小値を用いたDP ABC334-F

リストの途中に要素を入れたり出したりしたい

双方向リストを用いる

# 双方向リストとは直感的には、各要素が「自分の前の要素」と「自分の後の要素」の情報のみを持つデータ構造のことです。
# https://atcoder.jp/contests/abc344/editorial/9487

from collections import defaultdict

INF = pow(10,10)

N = int(input())
A = [int(e) for e in input().split()]

prev_value = dict()     # prev_value[i] : iの前に入っている値
next_value = dict()     # next_value[i] : iの次に入っている値

prev_value[0] = 0       # 入る値と重複しないようにすること。問題文より今回は入る値は1～なので、OK
next_value[0] = A[0]
prev_value[INF] = A[-1]
next_value[INF] = INF   # 入る値と重複しないようにすること。問題文より今回は入る値は～10**9なので、OK

# 初期状態の前後の要素を調べてまとめる
for i in range(N):
    if i == 0:
        prev_value[A[i]] = 0
        if N != 1:
            next_value[A[i]] = A[i+1]
        elif N == 1:
            next_value[A[i]] = INF
    elif i == N-1:
        prev_value[A[i]] = A[i-1]
        next_value[A[i]] = INF
    else:
        prev_value[A[i]] = A[i-1]
        next_value[A[i]] = A[i+1]

# クエリを処理する
Q = int(input())
for _ in range(Q):
    query = [int(e) for e in input().split()]

    if query[0] == 1:
        # xの直後にyを挿入する。
        x,y = query[1],query[2]
        # x -> z が x -> y -> z となる。
        z = next_value[x]
        next_value[x] = y
        prev_value[y] = x        
        next_value[y] = z
        prev_value[z] = y
        
    elif query[0] == 2:
        # xという値を削除する
        x = query[1]
        # a -> x -> b が a -> bとなる。
        a = prev_value[x]
        b = next_value[x]
        prev_value[b] = a        
        next_value[a] = b

A = list()
v = 0
while v != INF:
    if next_value[v] == INF:
        break
    else:
        A.append(next_value[v])
        v = next_value[v]

print(*A)

リストを処理して条件を満たす～を求めたい

条件を満たす最長連続部分列を求めたい

→ 左端から見ていってdequeで処理する
問題例 : ABC124-D「Handstand」 → 回答

条件を満たす順列を作るための最低操作回数を求めたい

→ 左端からdequeで処理しながら確定していく。
問題例 : AGC058-A「Make it Zigzag」 → 回答

文字列を取り扱いたい

最長共通部分列(LCS)を求めたい

e.g. 鉄則A20-LCS

DPする

S = list(input())
T = list(input())
M = len(S)
N = len(T)
S = ["#"] + S
T = ["#"] + T

lcs = [[0 for _ in range(N+1)] for _ in range(M+1)]
# lcs[i][j] : Sのi文字目,Tのj文字目まで見たときの最長共通部分列 

for i in range(1,M+1):
    for j in range(1,N+1):
        if S[i] == T[j]:
            lcs[i][j] = lcs[i-1][j-1] + 1
        else:
            lcs[i][j] = max(lcs[i-1][j],lcs[i][j-1])

print(lcs[M][N])

最長共通接頭辞(LCP)を求めたい

e.g. ABC287-E

先頭から一致確認をして、一致している集団毎にどんどん分割していく。

from collections import defaultdict

N = int(input())

S_list = list()
for i in range(N):
    S = str(input())
    S_list.append(S)

ans = defaultdict(int)
stack = list()
stack.append((0,S_list))

while len(stack) != 0:
    index,word_list = stack.pop(-1)

    # word_list内の単語をindexの文字毎にグループ分けする
    devided = defaultdict(list)

    for word in word_list:
        if len(word) > index:
            devided[word[index]].append(word)
        else:
            # LCPは短い方の文字数を越えないので、越えていたら終了
            ans[word] = index 

    for k,v in devided.items():
        if len(v) == 1:
            # indexの文字が一致するものはないので終了
            ans[v[-1]] = index
        else:
            # indexの文字が一致するものがあれば、さらに分割する
            stack.append((index+1,v))

for S in S_list:
    print(ans[S])

少しずつ変化する文字列が一致しているかを確認したい → ローリングハッシュ

数・数列を取り扱いたい

Xの倍数になる整数の組み合わせ = Xを素因数分解。整数の素因数分解で、Xの要素が全部含まれていればOK

切り上げ・切り下げしたい

math.ceilとかだと16桁程度くらいしか精度担保されないのでこうする

denominator,numerator = map(int, input().split())   # 分子,分母
ceiled = int(-(-denominator//numerator))    # 切り上げ：ceil(denominator/numerator)
floored = int(denominator//numerator)       # 切り捨て：floor(denominator/numerator)

桁数を求めたい

10で割れる回数を求めればよい

def get_number_of_digit(n):
    number_of_digit = 0

    if n == 0:
        number_of_digit = 1
    else:
        while n != 0:
            number_of_digit += 1
            n = n // 10

    return number_of_digit

print(get_number_of_digit(1234567))  # 7
print(get_number_of_digit(0))        # 1

log(10)で求める手もあるが、巨大な数とかで誤差が出たはず。

a^-1 % MOD を求めたい

フェルマーの小定理を利用するver

# https://twitter.com/kyopro_friends/status/1647460499263721478
# MODの世界ではフェルマーの小定理から P^(MOD-1)=1が成り立ってて、両辺をPで割るとP^(MOD-2)=1/Pになる
MOD_BY = 998244353
inv_a = pow(a, MOD_BY - 2,MOD_BY)

MODの逆元

MOD_BY = 998244353

# https://tex2e.github.io/blog/crypto/modular-mul-inverse
def xgcd(a, b):
    x0, y0, x1, y1 = 1, 0, 0, 1
    while b != 0:
        q, a, b = a // b, b, a % b
        x0, x1 = x1, x0 - q * x1
        y0, y1 = y1, y0 - q * y1
    return a, x0, y0

def modinv(a):
    # (1/a) % MOD_BY を求める
    m = MOD_BY
    g, x, y = xgcd(a, m)
    if g != 1:
        raise Exception('modular inverse does not exist')
    else:
        return x % m

n!や1/(n!)のMOD を求めたい

フェルマーの小定理を利用する

# e.g. アルゴリズムと数学　演習問題集 051 - Combination Hard
#      https://atcoder.jp/contests/math-and-algorithm/tasks/math_and_algorithm_ar

#import pypyjit
#pypyjit.set_param('max_unroll_recursion=0')
# ↑でPyPyの再帰が速くなるらしい https://qiita.com/shoji9x9/items/e7d19bd6f54e960f46be
# ★★★Pythonでの提出時はコメントアウトすること★★★

# 大きい数の場合は、デフォルトの再帰限界を超えるので拡張しておく
import sys
sys.setrecursionlimit(1000000)

MOD_BY = 1000000007		# 数値は問題文によって変えること!

def calc_mod_factorial(n):
    # n!をMODしたものを求める
    if n in mod_factorial:
        pass
    elif n == 1:
        mod_factorial[1] = 1
    else:
        mod_factorial[n] = n * calc_mod_factorial(n-1) % MOD_BY

    return mod_factorial[n]

def calc_inv_mod_factorial(n):
    # 1/(n!)をMODしたものを求める
    # = (1/n)%MOD * (1/(n-1))%MOD * ･･･
    # それぞれの項はフェルマーの小定理から求められる
    if n in inv_mod_factorial:
        pass
    elif n == 1:
        inv_mod_factorial[1] = pow(1, MOD_BY - 2,MOD_BY)
    else:
        inv_n = pow(n, MOD_BY - 2,MOD_BY)
        inv_mod_factorial[n] = inv_n * calc_inv_mod_factorial(n-1) % MOD_BY
    
    return inv_mod_factorial[n]

mod_factorial = dict()		# 複数個の値で求める場合は再利用を効かせるために記録
inv_mod_factorial = dict()	# 〃

X,Y = map(int, input().split())

mod_factorial = dict()
inv_mod_factorial = dict()

# 全ての移動の並べ方の通りは (X+Y)!
# 縦移動と横移動はそれぞれの中で可換なので X! * Y!で割ればOK
numerator_MOD = calc_mod_factorial(X+Y)
denominator_MOD = calc_inv_mod_factorial(X) * calc_inv_mod_factorial(Y) % MOD_BY
print(numerator_MOD * denominator_MOD % MOD_BY)

約数を列挙したい

約数列挙

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]

素数を列挙したい

素数列挙 : 3*(10**6)までで

pn_list = [2]

for q in range(3,3*(10**6) + 1):
    for a in pn_list:
        if q%a == 0:
            # qは素数でない
            break
        elif a > int(math.sqrt(q)):
            # qは素数
            pn_list.append(q)
            break

素数か判定したい

素数判定

import math

def is_prime(n):
    # https://qiita.com/srtk86/items/874639e361917e5016d4
    if n == 1: return False

    for k in range(2, int(math.sqrt(n)) + 1):
        if n % k == 0:
            return False

    return True

print(is_prime(100000000003))   # True 初期化含め60ms程度掛かる。

素因数分解したい

素因数分解

from collections import defaultdict

def factorize(n):
    b = 2
    fct = defaultdict(int)
    while b * b <= n:
        while n % b == 0:
            n //= b
            fct[b] += 1
        b = b + 1
        
    if n > 1:
        fct[n] += 1
    
    return fct

巨大な数の時はポラード・ロー素因数分解法を用いる。

数字の積が平方数か否かを判定したい

各数を素因数分解→各数毎に要素数を足し合わせ。全要素の個数ともに偶数ならOK　e.g.ABC342-D

from collections import defaultdict

N = int(input())
A = [int(e) for e in input().split()]


# Aを平方の積の積 * そうでない数 の表現に直して、まとめる
square_resolved = defaultdict(int)  # key:各数の平方の積で表すことのできない分,value:その個数
# ex 12 = pow(2,2) * 3　→ key:3
# ex 36 = pow(2,2) * pow(3,2) * 1　→ key:1

for i in range(N):
    j = 2
    Ai = A[i]
    while Ai >= pow(j,2):
        while Ai % pow(j,2) == 0:
            Ai = Ai // pow(j,2)

        j += 1
    
    square_resolved[Ai] += 1


# まとめた結果を利用して答えを出す
ans = 0

for k,v in square_resolved.items():
    if k == 0:
        # keyが0の場合は特殊ケース
        # Ai=0ならAjはどんな数でもAiAj=0で平方数になる
        ans += v * (N-1) - v * (v-1) // 2   
        # v * N             ･･･ Ai = 0のとき、その左側とその右側の全組み合わせが平方数になる
        # v * (v-1) // 2　  ･･･　0同士の組み合わせは二重カウントされているので、その分を除く
    elif k != 0:
        ans += v * (v-1) // 2
        # 各数の平方の積で表すことのできない分の中から2つを選ぶ組み合わせの通り数

print(ans)

数字が平方数かどうかの判定は単純にルートを取って、整数になればOK

最小公倍数/最大公約数を求めたい

lcm & gcd

import math

X,Y = map(int, input().split())
gcd_XY = math.gcd(X,Y)              # 最大公約数
lcm_XY = X * Y // math.gcd(X,Y)     # 最小公倍数

3つ以上の数に対しての最小公倍数

from collections import defaultdict

def factorize(n):
    # nを素因数分解する
    b = 2
    fct = defaultdict(int)
    while b * b <= n:
        while n % b == 0:
            n //= b
            fct[b] += 1
        b = b + 1
        
    if n > 1:
        fct[n] += 1
    
    return fct

def get_lcm(A):
    # リストで与えられた数字のlcm(最小公倍数)を求める
    ans_dict = defaultdict(int)
    for i in range(N):
        for k,v in factorize(A[i]).items():
            ans_dict[k] = max(ans_dict[k],v)

    lcm = 1
    for k,v in ans_dict.items():
        lcm *= pow(k,v)
    
    return lcm

N = int(input())
A = [int(e) for e in input().split()]

print(get_lcm(A))

3つ以上の数に対しての最大公約数

# リスト中の最小の数に対して約数全列挙している。
# 約数を最大側から順に処理すれば、不要な列挙を削れるので、高速化できるはず

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]

def get_gcd(A):
    # リストで与えられた数字のlcm(最大公約数)を求める
    A.sort(reverse=False)
    len_A = len(A)
    divisors = make_divisors(A[0])
    
    while len(divisors) != 0:
        remain_maxi_divisor = divisors.pop(-1)
        
        for i in range(1,len_A):
            if A[i] % remain_maxi_divisor != 0:
                break
        else:
            return remain_maxi_divisor

N = int(input())
A = [int(e) for e in input().split()]

print(get_gcd(A))

数字の桁を操作したい

整数の桁操作

def ope_number_digit(N,K,P):
    # 整数Nの左からK桁目をPに書き換える
    N -= get_number_digit(N,K) * pow(10,K-1)
    N += P * pow(10,K-1)
    return N

def get_number_digit(N,K):
    # 整数Nの左からK桁目を得る
    return N // pow(10,K-1) - N // pow(10,K) * 10

小数を含む式の結果を高精度に計算したい

Decimalモジュールを使う

from decimal import Decimal, getcontext
getcontext().prec = 12  # 必要な精度に変えること

a,b,c = map(Decimal, input().split())
maxi = Decimal('2.0')
mini = Decimal('1.0')

for i in range(50):
    x = (maxi+mini)/2
    
    if a * pow(x,5) + b * x + c >= 0:
        #midでOKだったので、OKとなる最小値はmidより小さい値である
        maxi = x
    else:
        #midでNGだったので、OKとなる最小値はmidより大きい値である
        mini = x

print(x)

他に

十分に大きい値を掛ける
小数を取り除いた形で読み込み、その分掛けたり割ったりする

方法などもある
10**-6程度の精度であれば、小数*小数となるような計算を避ければ、floatを用いるだけでOK e.g.ABC327-E

MEX(数列に含まれない最小の非負整数)を求めたい

数列に含まれない数を平方分割を利用したSorted (Multi) Setで管理する。

Sorted (Multi) Setを用いてMEXを求める e.g.ABC330E

# Sorted Multi Set
# 本体   →  https://github.com/tatyam-prime/SortedSet/blob/main/SortedMultiset.py
# 使用法 →  https://github.com/tatyam-prime/SortedSet
# 解説   →  https://qiita.com/tatyam/items/492c70ac4c955c055602
import math
from bisect import bisect_left, bisect_right
from typing import Generic, Iterable, Iterator, List, Tuple, TypeVar, Optional
T = TypeVar('T')

class SortedMultiset(Generic[T]):
    BUCKET_RATIO = 50
    REBUILD_RATIO = 170

    def _build(self, a: Optional[List[T]] = None) -> None:
        "Evenly divide `a` into buckets."
        if a is None: a = list(self)
        size = len(a)
        bucket_size = int(math.ceil(math.sqrt(size / self.BUCKET_RATIO)))
        self.a = [a[size * i // bucket_size : size * (i + 1) // bucket_size] for i in range(bucket_size)]
    
    def __init__(self, a: Iterable[T] = []) -> None:
        "Make a new SortedMultiset from iterable. / O(N) if sorted / O(N log N)"
        a = list(a)
        self.size = len(a)
        if not all(a[i] <= a[i + 1] for i in range(len(a) - 1)):
            a = sorted(a)
        self._build(a)

    def __iter__(self) -> Iterator[T]:
        for i in self.a:
            for j in i: yield j

    def __reversed__(self) -> Iterator[T]:
        for i in reversed(self.a):
            for j in reversed(i): yield j
    
    def __eq__(self, other) -> bool:
        return list(self) == list(other)
    
    def __len__(self) -> int:
        return self.size
    
    def __repr__(self) -> str:
        return "SortedMultiset" + str(self.a)
    
    def __str__(self) -> str:
        s = str(list(self))
        return "{" + s[1 : len(s) - 1] + "}"

    def _position(self, x: T) -> Tuple[List[T], int]:
        "Find the bucket and position which x should be inserted. self must not be empty."
        for a in self.a:
            if x <= a[-1]: break
        return (a, bisect_left(a, x))

    def __contains__(self, x: T) -> bool:
        if self.size == 0: return False
        a, i = self._position(x)
        return i != len(a) and a[i] == x

    def count(self, x: T) -> int:
        "Count the number of x."
        return self.index_right(x) - self.index(x)

    def add(self, x: T) -> None:
        "Add an element. / O(√N)"
        if self.size == 0:
            self.a = [[x]]
            self.size = 1
            return
        a, i = self._position(x)
        a.insert(i, x)
        self.size += 1
        if len(a) > len(self.a) * self.REBUILD_RATIO:
            self._build()
    
    def _pop(self, a: List[T], i: int) -> T:
        ans = a.pop(i)
        self.size -= 1
        if not a: self._build()
        return ans

    def discard(self, x: T) -> bool:
        "Remove an element and return True if removed. / O(√N)"
        if self.size == 0: return False
        a, i = self._position(x)
        if i == len(a) or a[i] != x: return False
        self._pop(a, i)
        return True

    def lt(self, x: T) -> Optional[T]:
        "Find the largest element < x, or None if it doesn't exist."
        for a in reversed(self.a):
            if a[0] < x:
                return a[bisect_left(a, x) - 1]

    def le(self, x: T) -> Optional[T]:
        "Find the largest element <= x, or None if it doesn't exist."
        for a in reversed(self.a):
            if a[0] <= x:
                return a[bisect_right(a, x) - 1]

    def gt(self, x: T) -> Optional[T]:
        "Find the smallest element > x, or None if it doesn't exist."
        for a in self.a:
            if a[-1] > x:
                return a[bisect_right(a, x)]

    def ge(self, x: T) -> Optional[T]:
        "Find the smallest element >= x, or None if it doesn't exist."
        for a in self.a:
            if a[-1] >= x:
                return a[bisect_left(a, x)]
    
    def __getitem__(self, i: int) -> T:
        "Return the i-th element.  (i番目の要素の値を返す)"
        if i < 0:
            for a in reversed(self.a):
                i += len(a)
                if i >= 0: return a[i]
        else:
            for a in self.a:
                if i < len(a): return a[i]
                i -= len(a)
        raise IndexError
    
    def pop(self, i: int = -1) -> T:
        "Pop and return the i-th element. (i番目の要素を削除して値を返す)"
        if i < 0:
            for a in reversed(self.a):
                i += len(a)
                if i >= 0: return self._pop(a, i)
        else:
            for a in self.a:
                if i < len(a): return self._pop(a, i)
                i -= len(a)
        raise IndexError

    def index(self, x: T) -> int:
        "Count the number of elements < x. (x未満の要素個の個数を調べる)"
        ans = 0
        for a in self.a:
            if a[-1] >= x:
                return ans + bisect_left(a, x)
            ans += len(a)
        return ans

    def index_right(self, x: T) -> int:
        "Count the number of elements <= x. (x以下の要素の個数を調べる)"
        ans = 0
        for a in self.a:
            if a[-1] > x:
                return ans + bisect_right(a, x)
            ans += len(a)
        return ans

N,Q = map(int, input().split())
A = [0] + [int(e) for e in input().split()]

# mexは最悪パターンでも,A=[0,1,...,N-1]でNとなる。
# ので、N以上の値はすべてNとして扱ってOK

# 各数が何個数列内にあるかを求める
cnt = [0 for _ in range(N+1)]   # mexの性質より、NまででOK
for i in range(1,N+1):
    if A[i] >= N:
        A[i] = N
    cnt[A[i]] += 1

# 数列内にない数字を求める...その中の最小値=mex
non_used_int = SortedMultiset([])
for i in range(N+1):
    if cnt[i] == 0:
        non_used_int.add(i)

# クエリを処理する
for _ in range(Q):
    i,x = map(int, input().split())
    x = min(x,N)    # mexの性質より、大きい値の時はNで切る。

    cnt[A[i]] -= 1
    if cnt[A[i]] == 0:
        # A内の個数が0になったので、"ない数字"に加える
        non_used_int.add(A[i])

    cnt[x] += 1
    non_used_int.discard(x) # A内にあるので、"ない数字"から除去(discardなので、そもそもない場合は何も起こらない)

    A[i] = x

    mex = non_used_int.ge(0)
    print(mex)

数列などの和を求めたい

$ L+(L+1)+(L+2)+...+(H-1)+H = (L+H)*(H-L+1)/2 $　... 等差数列のガウスの足し算
$ 1^2 + 2^2 + ... + N^2 = N(N-1)(N-2)/6 $
$ 2^0 + 2^1 + ... + 2^{N-1}= 2^N - 1 $
$ a+ ar + ar^1 + ar^2 + ... + ar^n = a(r^n - 1)/(r - 1)$ ... 等比数列の和の公式
MODを求める問題のときは分子をフェルマーの小定理を用いて掛け算に直して計算する
$ a+ (a+d) + (a+2d) + ... + (a + (n-1)d) = n(2a + (n-1)d)/2$ ... 等差数列の和の公式
$ 1+ p + p^1 + p^2 + ... = 1/(1-p) $ (ただし 0<p<1)
$ 1+ p + p^1 + p^2 + ... + p^n $ (ただし 1>=p) ... 再帰関数に落とし込む

一般項を求めて再帰に落とし込む e.g.ABC298-E

n = int(input())
m = 1000000007
p = 4

# p**0 ～ p**n の和を求める

# p**0 ～ p**(i-1) までの和をS[i]と表すとしたとき
# S[i+1] = p**0 + p**1 + … + p**(i-1) + p**i
#        = p**0 + p * S[i]
#        = 1 + p * S[i]
# S[2*i] = p**0 + … + p**(i-1) + p**(i) + … + p**(2*i - 1) 
#        = S[i]                + p**(i) * S[i]
#        = S[i] * (1 + p**i)
# となる。
# これを利用し再帰する。

def rec(i):
    
    if i == 0:
        S = 0
    else:
        if i%2 == 0:
            S = rec(i//2) * (1 + pow(p,i//2,m))
        else:
            S = 1 + p * rec(i-1)

    return S % m

print(rec(n+1)%m)

和の記号Σの性質を利用する

\begin{array}{ll}
\displaystyle \sum_{i=1}^n (a_{i} + b_{i}) = \sum_{i=1}^n a_{i} + \sum_{i=1}^n b_{i} \\
\displaystyle \sum_{i=1}^n ca_{i} = c \sum_{i=1}^n a_{i} \\
\displaystyle \sum_{i=1}^n c = c * n \\
\displaystyle \sum_{1\leqq i < j \leqq n} c = c * n*(n-1)/2 \\
\displaystyle \sum_{i=1}^n \sum_{j=1}^n A_{ij} = \sum_{j=1}^n \sum_{i=1}^n A_{ij} \\
\displaystyle \sum_{1\leqq i < j \leqq n} A_{ij} = \sum_{i=1}^n \sum_{j=i+1}^n A_{ij} 
\end{array}

(4段目補足) 1~nの中から異なるi,jの組み合わせの個数を掛け合わせることになるので。
i = 1~N でなく i = N～1 と考えると有効な場合がある

指数を用いた計算を簡易化したい(指数法則)

あとで書く

図形の～を求めたい

三角関数を取り扱いたい

三角関数

# 例:半径と角度からxy座標を求める
import math

r,deg = map(int, input().split())   # 2,60

rad = math.radians(deg)
x = r * math.cos(rad)
y = r * math.sin(rad)

print(rad,x,y)  #1.0471975511965976 1.0000000000000002 1.7320508075688772

# 【floatで誤差が出るので.iscloseを使うこと】
print(x == 1 and y == math.sqrt(3),1,math.sqrt(3))                                      # False 1 1.7320508075688772
print(math.isclose(x,1,abs_tol=1e-10) and math.isclose(y,math.sqrt(3),abs_tol=1e-10))   # True

辺AB,BCのなす角を求めたい

なす角を求める

import math

def get_deg(Ax,Ay,Bx,By,Cx,Cy):
    # 辺ABと辺BCのなす角(0～180)を求める
    # cf.https://w3e.kanazawa-it.ac.jp/math/category/vector/henkan-tex.cgi?target=/math/category/vector/naiseki-wo-fukumu-kihonsiki.html
    ax = Ax - Bx
    ay = Ay - By
    bx = Cx - Bx
    by = Cy - By

    if ax**2 + ay**2 != 0 and bx**2 + by**2 != 0:
        cos_theta = (ax*bx + ay*by)/(math.sqrt(ax**2 + ay**2) * math.sqrt(bx**2 + by**2))
        return math.degrees(math.acos(cos_theta))
    else:
        return -1

3点の座標から三角形の面積を求めたい

3点の座標から三角形の面積を求める

def get_triangle_size_from_three_points(xa,ya,xb,yb,xc,yc):
    # https://mathwords.net/x1y2hikux2y1#i-2
    return abs((xa-xc)*(yb-yc)-(xb-xc)*(ya-yc)) / 2
    # /2 で誤差が出るので、必要に応じてmath.iscloseを使うこと

円周上の弦2本が交差するか否か

弦を始点・終点を数直線上で考えて、重なる部分があれば交差する

# e.g. パ研合宿2023 Day1-B「Cutting Circle」
# https://atcoder.jp/contests/pakencamp-2023-day1/tasks/pakencamp_2023_day1_b

N = int(input())
a,b = map(int, input().split())
c,d = map(int, input().split())

# 直線が交差するなら4分割、しないなら3分割

# 扱いやすいように大小を整理
if a > b:
    a,b = b,a
if c > d:
    c,d = d,c

# 交差の有無を判定
if c < a:
    if a < d < b:
        is_clossing = True
    else:
        is_clossing = False
elif c == a or c == b:
    is_clossing = False     # 2辺が接している。c=a かつ d=bなら 辺が重なっている
elif a < c < b:
    if d <= b:
        is_clossing = False # d=bのときは2辺が接している
    else:
        is_clossing = True
elif b < c:
    is_clossing = False     # c<d　なので　b<=d　が確定。

if is_clossing:
    print(4)
else:
    print(3)

マンハッタン距離が関係する問題を解きたい

各次元ごとの操作を独立に行うことが出来る場合があることを利用する
e.g. 全点間のマンハッタン距離の和を求めるのを各軸毎に行って和を取る
マンハッタン距離の和の最小化は各軸に分け、それぞれの真ん中の値にすればOK
各座標のx,yに対して z = x+y,w = x-y と置いたとき、マンハッタン距離はmax(|zA−zB|,|wA−wB|)
∴多点でのマンハッタン距離の最大値は max(max(z)-min(z),max(w)-min(w))

行列を取り扱いたい

行列の入力を読み込みたい

→ H行W列の二次元形式な入力を読み込みたい

行列を回転/反転/転置したい

行列の回転/反転/転置

def print_2d_matrix(M):
    # 表示用関数
    W = len(M[0])
    H = len(M)

    for i in range(H):
        print(*M[i])
    
    print("")
    
#【!】回転・転置による行列幅の入れ替えに注意せよ!!
A = [[i+j*3 for i in range(3)] for j in range(5)]
A = list(map(list, zip(*A[::-1])))              #90度右回転
A = list(zip(*A))[::-1]                         #90度左回転
A = [A[i][::-1] for i in range(len(A))][::-1]   #180度回転
A = list(zip(*A))                               #転置(行と列の入れ替え)
A = A[::-1]                                     #上下反転
A = [A[i][::-1] for i in range(len(A))]         #左右反転

bit・二進数を取り扱いたい

bit演算を取り扱いたいときはこれ。

bit・二進数の基本操作

# 下記のnは一番右のフラグを1番目としているので注意。
# "右"からn桁目のフラグたってるかの判定
if target & (1<<n-1):  # == 1 はつけないこと!!
    print("右から",n,"番目のbitはON")
else:
    print("右から",n,"番目のbitはOFF")

#"右"からn桁目のフラグをon
target = target | (1<<n-1)
#"右"からn桁目のフラグをoff
target = target & ~(1<<n-1)

# "左"からn桁目～で考えるときは、フラグ全ONを(2**N)-1で表せるとき下記となる
m = N - n
if target & (1<<m):  # == 1 はつけないこと!!
    print("左から",n,"番目のbitはON")
else:
    print("左から",n,"番目のbitはOFF")

#左からn桁目のフラグをon
target = target | (1<<m)
#左からn桁目のフラグをoff
target = target & ~(1<<m)


# bit毎のXOR
X = A^B        # A^B = X ⇔ A = B^X ⇔ B = A^X 。もちろん A^B = B^A

二進数を最大化したい場合はより上位桁を立てる、最小化したい場合はより上位桁を折る
e.g. 0b1000 > 0b0111 (8 > 7) ... ABC281-F

集合同士の結合して、同一集合か判定したり、集合の大きさを求めたい

重み無の場合

Union Findを使う

# Union Find ###################################################################
# https://note.nkmk.me/python-union-find/
import sys
from collections import defaultdict

sys.setrecursionlimit(1000000)

class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n     #各要素の親番号が入る。根の場合は-(グループ要素数)が入る

    def find(self, x):
        #要素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が属するグループと要素yが属するグループとを併合する

        x_root = self.find(x)    # 要素xが属するグループの根
        y_root = self.find(y)    # 要素yが属するグループの根

        if x_root == y_root:
            # 根が同じなら何もしない
            return

        # 併合元/先 (どちらを根にするか)の変更処理。必要に応じて下記A～Cから選ぶ
        # 処理しない場合は、要素xが属するグループに要素yが属するグループが取り込まれる形となる

        # A.グループのサイズが大きい方を根にする
        if self.parents[x_root] > self.parents[y_root]:
            x_root, y_root = y_root, x_root

        '''
        #B.グループの根の番号が大きい方を根にする
        if x_root < y_root:
            x_root, y_root = y_root, x_root

        #C.グループの根の番号が小さい方を根にする
        if x_root > y_root:
            x_root, y_root = y_root, x_root
        '''

        #併合する
        self.parents[x_root] += self.parents[y_root]    #要素が根の場合は、-(グループの要素数)になるように処理
        self.parents[y_root] = x_root                   #yの親要素をxに書き換え

    def size(self, x):
        #要素xが属するグループのサイズ（要素数）を返す
        return -self.parents[self.find(x)]

    def same(self, x, y):
        #要素x, yが同じグループに属するかどうかを返す
        return self.find(x) == self.find(y)

    def members(self, x):
        #要素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):
        #{ルート要素: [そのグループに含まれる要素のリスト], ...}のdefaultdictを返す
        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())
################################################################################

uf = UnionFind(N+1)     # union find。要素0～Nまでが使えるようになる
uf.find(x)              # 要素xが属するグループの根を返す
uf.union(x,y)           # 要素xが属するグループと要素yが属するグループとを併合する
uf.same(x,y)            # 要素x, yが同じグループに属するかどうかを返す
uf.size(x)              # 要素xが属するグループのサイズ（要素数）を返す
uf.group_count()        # グループの数を得る
group_members =  uf.all_group_members()
for root,members in group_members.items():
    print(root,members)

縦N横Mの二次元座標(i,j) (0<=i<N,0<=j<M)で行いたいときはUnionFind(N*M)で初期化して、 i*M + j で1次元に直してやればOK
集合内での最大の数を求める問題では、別途要素毎に値を保持するリストを作って、根の要素が集合内での最大の数を持つようにすればOK
結合する際に、大きい方に小さい方を結合することと集合の最大値を求める際に根まで辿るのを忘れないこと。

重みありの場合

重み付きUnion Findを使う e.g.ABC328-F

# import pypyjit
# pypyjit.set_param('max_unroll_recursion=0')
# ↑でPyPyの再帰が速くなるらしい https://qiita.com/shoji9x9/items/e7d19bd6f54e960f46be
# ★★★Pythonでの提出時はコメントアウトすること★★★

# Weighted Union Find ###################################################################
# Union Find Original - https://note.nkmk.me/python-union-find/
# Weighted化          - https://at274.hatenablog.com/entry/2018/02/03/140504

from logging import root
import sys
from collections import defaultdict

sys.setrecursionlimit(1000000)

class WeightedUnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [i for i in range(n)]    # 各要素の親番号が入る
        self.rank = [0] * n                     # 木の高さ
        self.weight = [0] * n                   # 親への重みを管理
        self.size = [1] * n                     # 木のサイズ
        
    def find(self, x):
        #要素xが属するグループの根を返す
        if self.parents[x] == x:
            return x
        else:
            #根でないので、さらに辿る
            y = self.find(self.parents[x])                  # 最上位の親(根)が入る
            self.weight[x] += self.weight[self.parents[x]]  # 途中経路の親への重みを書き換え
            self.parents[x] = y                             # 途中経路の親の値を書き換え
            return y

    def union(self, x, y, w):
        # 要素xが属するグループと要素yが属するグループとを併合する
        # ただし、xからyの重みはw。(weight[y] = weight[x] + w)

        x_root = self.find(x)    # 要素xが属するグループの根
        y_root = self.find(y)    # 要素yが属するグループの根

        if x_root == y_root:
            # 根が同じなら重みの矛盾がないかどうかを返す
            if self.diff(x,y) == w:
                return True
            else:
                return False

        if self.rank[x_root] < self.rank[y_root]:
            # xの木の高さ < yの木の高さのとき処理
            self.parents[x_root] = y_root                               # 低い方の根から高い方の根に併合する
            self.weight[x_root] = w - self.weight[x] + self.weight[y]   # x_root→y_rootの重みを定める。x→x_root分を逆流するので-、x→yは順方向なので+,y→y_root分も順方向なので+。
            self.size[y_root] += self.size[x_root]                      # 併合した分、木のサイズを追加
        else:
            # xの木の高さ >= yの木の高さのとき処理
            self.parents[y_root] = x_root                               # 低い方の根から高い方の根に併合する
            self.weight[y_root] = -w - self.weight[y] + self.weight[x]  # y_root→x_rootの値を定める。y→y_rootは逆方向なので-、x→yも逆方向なので-,x→x_rootは順方向なので+
            self.size[x_root] += self.size[y_root]                      # 併合した分、木のサイズを追加
            
            if self.rank[x_root] == self.rank[y_root]:
                self.rank[x_root] += 1
        
        return True

# これから、ちゃんと動作するかちょっと怪しい
    def size(self, x):
        #要素xが属するグループのサイズ（要素数）を返す
        return self.size[self.find(x)]

    def same(self, x, y):
        #要素x, yが同じグループに属するかどうかを返す
        return self.find(x) == self.find(y)

    def members(self, x):
        #要素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 == x]

    def group_count(self):
        #グループの数を返す
        return len(self.roots())

    def all_group_members(self):
        #{ルート要素: [そのグループに含まれる要素のリスト], ...}のdefaultdictを返す
        group_members = defaultdict(list)
        for member in range(self.n):
            group_members[self.find(member)].append(member)
        return group_members
    
    def diff(self, x,y):
        # x→yのコストを求める
        return self.weight[x] - self.weight[y]

    def __str__(self):
        return '\n'.join(f'{r}: {m}' for r, m in self.all_group_members().items())
# ここまで、ちゃんと動作するかちょっと怪しい

################################################################################

N,Q = map(int, input().split())

w_uf = WeightedUnionFind(N+1)     # 重み付き union find。要素0～Nまでが使えるようになる
ans = list()

for i in range(Q):
    a,b,d = map(int, input().split())

    if w_uf.union(a,b,d):
        ans.append(i+1)

print(*ans)

グラフの～/～なグラフを求めたい

グラフの最短経路を求めたい

最短経路を出力する場合は、ゴールからスタート側に辿ればOK

重みなし単一始点最短経路問題

幅優先探索(二次元)　

from collections import deque

# 入力部省略 H*Wの二次元配列A,開始座標(Si,Sj),終了座標(Gi,Gj)

# 幅優先探索に必要な配列を準備
ans = [[-1 for _ in range(W+2)] for _ in range(H+2)]            # マス(i,j)にたどり着いたときの演算結果
is_visited = [[False for i in range(W+2)] for _ in range(H+1)]  # マス(i,j)へ訪問済みか
search_deq = deque()                                            #探索候補一覧

# 開始点の処理
ans[Si][Sj] = 0
search_deq.append((Si,Sj))
is_visited[Si][Sj] = True

# 幅優先探索の実施
while len(search_deq)!=0:
    i,j = search_deq.popleft()

    for u,v in [(1,0),(-1,0),(0,1),(0,-1)]:     # 上下左右への移動の場合。問題によって変えること
        if 1 <= i+u <= H and 1 <= j+v <= W and A[i+u][j+v] != "#" and is_visited[i+u][j+v] == False:
            is_visited[i+u][j+v] = True
            ans[i+u][j+v] = ans[i][j] + 1       # 全マス間距離1の場合。問題によって変えること
            search_deq.append((i+u,j+v))

print(ans[Gi][Gj])

幅優先探索(グラフ)　

from collections import deque

N,M = map(int, input().split())             # N:頂点数 , M:辺の本数

edge = [list() for _ in range(N+1)]         # edge[i] : 頂点iと結ばれている頂点のlist
for i in range(M):
    A,B = map(int, input().split())
    edge[A].append(B)
    edge[B].append(A)

dist = [-1 for _ in range(N+1)]             # 1→頂点iまでの距離
prev_node = [-1 for _ in range(N+1)]        # 1→頂点iまでの最短経路での直前の点
is_visited = [False for i in range(N+1)]    # 頂点iに訪問済みか

search_deq = deque()                        #探索候補一覧

# 開始点の処理
dist[1] = 0
search_deq.append(1)
is_visited[1] = True

# 幅優先探索
while len(search_deq)!=0:
    i = search_deq.popleft()

    for j in edge[i]:
        if is_visited[j] == False:
            is_visited[j] = True
            dist[j] = dist[i] + 1           # 全マス間距離1の場合。問題によって変えること
            prev_node[j] = i

            search_deq.append(j)

for i in range(1,N+1):
    print(dist[i])

O(E)
無限に広がるグリッド上での探索では単純に行うとTLEする。制約を利用して範囲を抑え込む。 e.g.:障害物・ゴールの座標制限より±1範囲外は探索不要

深さ優先探索(二次元)　

H,W = map(int, input().split())
edge = [["$" for _ in range(W+2)]]
center = [["$"] + list(input()) + ["$"] for i in range(H)]
S = edge + center + edge

move_volume = [(1,0),(-1,0),(0,1),(0,-1)]                       # 周囲のマスへの移動量
is_visited = [[False for _ in range(W+2)] for _ in range(H+2)]  # マス(i,j)が訪問済みか
researched = [[0 for _ in range(W+2)] for _ in range(H+2)]      # マス(i,j)から移動できるマスの何番目まで調べたか

stack = list()              # 探索ルートを記録するstack

si,sj = 1,1                 # (si,sj)は探索開始の座標
stack.append((si,sj))
is_visited[si][sj] = True
visited_cnt = 1             # 訪問済みの"#"なセルのカウント

while len(stack) != 0:
    i,j = stack[-1]

    for k in range(researched[i][j],4):
        # 隣接マスへの遷移を試す

        researched[i][j] += 1

        u,v = move_volume[k]    # マスでない場合は、別途移動可能頂点をまとめておいて、そこから読み込む

        if is_visited[i+u][j+v] == False and S[i+u][j+v] == "#":
            # 条件を満たす未訪問隣接セルがあれば遷移する
            is_visited[i+u][j+v] = True
            visited_cnt += 1
            stack.append((i+u,j+v))

            # ★ゴールに到達したら終了のパターンならば、ここで判定する。

            break               # 深さ優先なので、深い方に遷移出来たら抜ける

    else:
        # 全隣接セルが遷移済みの場合の処理

        # ★全てを巡る遷移が出来るかみたいな問題ならここで判定する

        # 1手戻す
        i,j = stack.pop(-1)
        is_visited[i][j] = False
        researched[i][j] = 0
        visited_cnt -= 1

深さ優先探索(グラフ) e.g.DPまとめコンG

# 問題を解くために有向辺の逆で取り扱っているので、無向や正方向で使うときは書き換え要!!
N,M = map(int, input().split())

inv_edge = [list() for _ in range(N+1)] # inv_edge[i]:頂点i"へ"有効辺を張っている頂点
for _ in range(M):
    x,y = map(int, input().split())
    inv_edge[y].append(x)

is_visited = [False for _ in range(N+1)]    # 頂点iが訪問済みか
dist = [0 for _ in range(N+1)]              # 頂点iへの最長有向パス長さ
researched = [0 for _ in range(N+2)]        # 頂点iへ移動できる頂点の何番目まで調べたか

stack = list()                              # 探索ルートを記録するstack
ans = 0                                     # グラフ中の最長有向パス長

for i in range(1,N+1):
    if is_visited[i] == False:
        stack.append(i)
        is_visited[i] = True

        while len(stack) != 0:
            v = stack[-1]

            for k in range(researched[v],len(inv_edge[v])):
                u = inv_edge[v][k]  # vへ有向辺を張っている頂点
                researched[v] += 1

                if is_visited[u] == True:
                    # 当該頂点の最長有向パスが計算済みなので利用する
                    dist[v] = max(dist[v],dist[u] + 1)
                else:
                    # 当該頂点の最長有向パスが計算未なので求める
                    is_visited[u] = True
                    stack.append(u)
                    break   # 深さ優先なので、深い方に遷移出来たら抜ける

            else:
                # vへ有向辺を張っている頂点を全て調べ済みの場合の処理
                u = stack.pop(-1)
                
                if len(stack) != 0:
                    # 下流側の頂点に最長有向パスを伝える
                    v = stack[-1]
                    dist[v] = max(dist[v],dist[u] + 1)
                else:
                    # 最下流なので、グラフ全体の最長有向パスの更新を試みる
                    ans = max(ans,dist[u])

print(ans)

重みあり(負の重みなし)単一始点最短経路問題

ダイクストラ法

# https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_ek
import heapq
INF = pow(10,20)                        # 無限大とする値。距離や頂点数から考えて十分大きい値にすること

N,M = map(int, input().split())         # N:頂点数 M:入力辺数

edge = [set() for _ in range(N+1)]      # edge[i]:頂点iから出ている辺

for i in range(M):
    A,B,C = map(int, input().split())   # A:出発元頂点番号 B:行先頂点番号 C:重み
    edge[A].add((B,C))                  # (同一頂点に向かう辺は1本であることが問題文で保証されていない場合は、最短辺のみが残るように工夫すること)
    edge[B].add((A,C))                  # 無向グラフなので逆側にも入れる


dist = [INF for _ in range(N+1)]        # 始点から当該頂点までの最短距離。各頂点に対し無限大で初期化。
dist[1] = 0                             # 始点から始点の距離なので当然0
prev_node = [-1 for _ in range(N+1)]    # 始点から当該頂点までの最短距離時の直前頂点。各頂点に対し、prev(v)←無し。これを辿れば最短経路が求まる

hq = [(0,1)]                            # 優先度付きキューに(始点からの距離,始点の頂点番号)を追加
# ↑何を優先して探索させるかを十分に検討して、入れる要素の順番を決めること。間違えるとTLEする!!!

heapq.heapify(hq)

while len(hq)!=0:
    d,u = heapq.heappop(hq)             # 優先度付きキューから始点に近い側から頂点を取り出し
    if dist[u] == d:                    # (Q(v)←altではなく、更新毎に追加しているので、不要パターンがある。最短の場合のみ処理)
        for v,C in edge[u]:             # 取り出した頂点から出ている全ての辺に対して
            new_d = d + C               #     新規の距離を計算し
            if dist[v] > new_d:         #     短くなっていたら更新
                dist[v] = new_d
                prev_node[v] = u
                heapq.heappush(hq,(new_d,v))    #更新したところの接続点は短くなる可能性があるので再調査

# 最短経路を求める...ゴールから辿る
ans = [N]
v = N
while 1:
    v = prev_node[v]    # 上流の頂点を求める
    
    if v == -1:
        # それより上流に頂点が無いので終わる
        break
    else:
        ans.append(v)   # 上流の頂点を追加する
    
ans = ans[::-1]         # 下流から上流の順番になっているので、ひっくり返す
print(*ans)

重みあり(負の重みなし)単一始点最短経路問題のみならず、経路によって各頂点が取りうる値が変わる問題にも使用可能 e.g.ABC305E

重みあり全点間最短経路問題

ワーシャルフロイド法

INF = pow(10,20)

N,M = map(int, input().split())     # N:頂点数, M:辺数

dist = [[0 if j == i else INF for j in range(N+1)] for i in range(N+1)]   # dist[A][B]:A→Bの最小コスト

# 入力読み込み＋初期化
for i in range(M):
    A,B,C = map(int, input().split())
    dist[A][B] = min(dist[A][B],C)
    dist[B][A] = min(dist[A][B],C)  # 無向のときは逆方向も加える

# 本計算
for k in range(1,N+1):
    for i in range(1,N+1):
        for j in range(1,N+1):
            if dist[i][j] > dist[i][k] + dist[k][j] and dist[i][k] != INF and dist[k][j] != INF:
                dist[i][j] = dist[i][k]+dist[k][j]

重みなし全点間最短経路問題

各頂点を始点として幅優先探索を行えばよい。

重みなし多点間最短距離問題

LCAを求めて頂点⇔LCAの距離を足せば行えばよい。

重みあり単一始点少辺数最短経路問題

DPする。 e.g. ABC291-F

各点訪問条件付き多点間最短距離問題

Bit DPする。

Bit DP e.g. ABC338-F

INF = pow(10,8) # ←問題に合わせて適正な値に変更すること!!!

N,M = map(int, input().split())     # N:頂点数, M:辺数

# １．全点間の距離を求めておく
# 今回はワーシャルフロイド法を用いる。幅優先探索などを用いる場合もある。
dist = [[0 if j == i else INF for j in range(N)] for i in range(N)]   # dist[A][B]:A→Bの最小コスト

# １-１．入力読み込み＋初期化
for i in range(M):
    U,V,W = map(int, input().split())
    U -= 1  # 頂点番号を0始まりにオフセット
    V -= 1  # 頂点番号を0始まりにオフセット
    dist[U][V] = min(dist[U][V],W)  # 複数辺あるときに最小重みの辺のみ残す
    # dist[V][U] = min(dist[V][U],W)  # 無向辺の場合は、こちらも行うこと


# １-２．本計算
for k in range(N):
    for i in range(N):
        for j in range(N):
            if dist[i][j] > dist[i][k] + dist[k][j] and dist[i][k] != INF and dist[k][j] != INF:
                dist[i][j] = dist[i][k]+dist[k][j]



# ２．bit DPで最良値を求める
# ２ー１．配列確保･･･影響を及ぼすものを列挙し、その分を確保する。
# 今回は訪問済みの頂点一覧と、今居る頂点
dp = [[INF for _ in range(N)] for _ in range(2**(N))]   # dp[i][j]:訪問済み頂点を二進数で表したものをi、今居る頂点がjのときの、重み累計最小

# ２ー２．配列初期化･･･初期/終了状態で値が定まっているものがあれば、入れておく。
# 今回は、周り始めでは重み累計が0であることが定まっている、
for i in range(N):
    next_bit_state = 0 | (1<<i)
    dp[next_bit_state][i] = 0

# ２ー３．DP
# 遷移を考える：ある状態 から/へ 変化できる状態を列挙する
# 今回は、今居る頂点から有向辺がある頂点へ遷移可能 → 配るDPをする。
for now_bit_state in range(1,2**N):
    # now_bit_state : 現在の訪問済み頂点を二進数で表したもの
    for now_pos in range(N):
        # now_pos : 今いる頂点
        if dp[now_bit_state][now_pos] != INF:
            # (dp[now_bit_state][now_pos] == INFでは、その状態をそもそも取りえないので、その先の遷移は考える必要なし)
            for next_pos in range(N):
                # next_pos : 次に行く頂点
                if dist[now_pos][next_pos] != INF:
                    next_bit_state = now_bit_state | (1<<next_pos)  # next_posに行った場合の訪問済み頂点を二進数で表したもの
                    dp[next_bit_state][next_pos] = min(dp[next_bit_state][next_pos],dp[now_bit_state][now_pos] + dist[now_pos][next_pos])

# (追加でbitを立てたものに配っている　→ 自分より値の大きいものしか変化しない → 後から配布元側の値が更新されることはない。)


# dpの結果で全点訪問済みのものの中から最小値を取り出して出力
ans = min(dp[2**N - 1])

if ans == INF:
    print("No")
else:
    print(ans)

巡回セールスマン問題。N <= 18 くらいのときによく使う。

各点選択条件付き最良選択問題

Bit DPする。

Bit DP e.g. ABC318-D

N = int(input())
D = [[0 for _ in range(i+1)] + [int(e) for e in input().split()] for i in range(N-1)]

# bit DPで最良値を求める
# １．配列確保･･･影響を及ぼすものを列挙し、その分を確保する。
# 今回は既に選んだ辺の端点に含まれている頂点一覧
dp = [0 for _ in range(2**N)]   # dp[i]:選択済み頂点のbit表現がiのときの重み最大値。1<=Dなので0で初期化すれば、十分。

# ２．配列初期化･･･初期/終了状態で値が定まっているものがあれば、入れておく。
# 今回は、辺を選んでいないときは0であることが定まっている
# dp[0] = 0     # 既に初期値を0で入れているので不要

# ３．DP
# 遷移を考える：ある状態 から/へ 変化できる状態を列挙する
# 今回は、未選択の頂点2点を追加で選ぶ場合であれば遷移可能 → 配るDPをする。
for now_bit_state in range(2**N):
    # now_bit_state : 現在の選択済み頂点を二進数で表したもの

    for i in range(N-1):
        # i : 選択する頂点1つめ
        if now_bit_state & (1<<i):  # == 1 はつけないこと!!
            # 既に頂点iは選択済みで、この場合の計算は不要なので、pass
            pass
        else:
            for j in range(i+1,N):
                # j : 選択する頂点2つめ
                if now_bit_state & (1<<j):  # == 1 はつけないこと!!
                    # 既に頂点jは選択済みで、この場合の計算は不要なので、pass
                    pass
                else:
                    # i,jを選択した場合のbit表現を求める
                    next_bit_state = now_bit_state  | (1<<i)     # 左からi+1番目のbitを立てる
                    next_bit_state = next_bit_state | (1<<j)    # 左からj+1番目のbitを立てる
                    # i-jを結ぶ辺を追加で選択した場合の重み総和を求める
                    dp[next_bit_state] = max(dp[next_bit_state],dp[now_bit_state] + D[i][j])

# (追加でbitを立てたものに配っている　→ 自分より値の大きいものしか変化しない → 後から配布元側の値が更新されることはない。)

# 解答を求め、出力
print(max(dp))

順序が関係ないときは、こちらでやる。

各点非訪問条件付き単一始点最短距離問題

スタートから各頂点への距離と各頂点からゴールまでの距離を求め、非訪問点を飛ばすようにする。 e.g. ABC291-F

グラフの最長経路を求めたい

任意の点を始点に幅優先探索をし、仮の最遠点を求める。そこからさらに幅優先探索することで求める。

e.g.ABC361E

from collections import deque

N = int(input())

edge = [list() for _ in range(N+1)]

for i in range(N-1):
    A,B,C = map(int, input().split())
    edge[A].append((B,C))
    edge[B].append((A,C))

# 幅優先探索で木の中の最長経路を求める


# まず1を始点とした場合の最遠頂点を求める
# 幅優先探索準備
dist = [-1 for _ in range(N+1)]             # 1→頂点iまでの距離
is_visited = [False for i in range(N+1)]    # 頂点iに訪問済みか
search_deq = deque()                        #探索候補一覧

# 開始点の処理
farthest_node_a = 1
dist[1] = 0
search_deq.append(1)
is_visited[1] = True

# 幅優先探索
while len(search_deq)!=0:
    i = search_deq.popleft()

    for j,k in edge[i]:
        if is_visited[j] == False:
            is_visited[j] = True
            dist[j] = dist[i] + k           # 全マス間距離1の場合。問題によって変えること

            search_deq.append(j)

            if dist[farthest_node_a] < dist[j]:
                farthest_node_a = j


# 上で求めた最遠点を始点として最遠頂点を求める → 木の中の最長経路が求まる
# 幅優先探索準備
dist = [-1 for _ in range(N+1)]             # 1→頂点iまでの距離
prev_node = [-1 for _ in range(N+1)]        # 1→頂点iまでの最短経路での直前の点
is_visited = [False for i in range(N+1)]    # 頂点iに訪問済みか
search_deq = deque()                        #探索候補一覧

# 開始点の処理
dist[farthest_node_a] = 0
search_deq.append(farthest_node_a)
is_visited[farthest_node_a] = True
farthest_node_b = farthest_node_a

# 幅優先探索
while len(search_deq)!=0:
    i = search_deq.popleft()

    for j,k in edge[i]:
        if is_visited[j] == False:
            is_visited[j] = True
            dist[j] = dist[i] + k           # 全マス間距離1の場合。問題によって変えること
            prev_node[j] = i

            search_deq.append(j)

            if dist[farthest_node_b] < dist[j]:
                farthest_node_b = j

# ここまでで最長経路(farthest_node_a～間はprev_nodeで求める～farthest_node_b)が求まった。
# 以後省略

木の頂点毎の～を求めたい

全方位木DP

# ある頂点を根にしたときの部分木の値を求めておいて、そこから全体を求める。
# その結果や部分木の結果を利用して、他の頂点での結果を求めていく。
# cf1: https://qiita.com/katsumata_yusuke/items/c17ca95602ac2a4c8322
# cf2: https://twitter.com/kyopro_friends/status/1442128148095664144

N = int(input())

edge = [list() for _ in range(N+1)]         # edge[i]:頂点iから辺を張っている頂点
for _ in range(N-1):
    u,v = map(int, input().split())
    edge[u].append(v)
    edge[v].append(u)

ans = [-1 for _ in range(N+1)]

# 根を1として、部分木での値を深さ優先探索で求める
is_visited = [False for _ in range(N+1)]    # 頂点iが訪問済みか
count = [0 for _ in range(N+1)]             # 頂点iより下流側の部分木内の頂点数
dist_sum = [0 for _ in range(N+1)]          # 頂点iより下流側の部分木内部のdist(i,j)の合計
researched = [0 for _ in range(N+2)]        # 頂点iから移動できる頂点の何番目まで調べたか

stack = list()                              # 探索ルートを記録するstack
stack.append(1)
is_visited[1] = True
count[1] = 1

while len(stack) != 0:
    u = stack[-1]

    for k in range(researched[u],len(edge[u])):
        v = edge[u][k]      # uから辺がある頂点
        researched[u] += 1

        if is_visited[v] == False:
            is_visited[v] = True
            count[v] = 1            # 自分自身の分
            stack.append(v)
            break   # 深さ優先なので、深い方に遷移出来たら抜ける

    else:
        # uから辺を張っている頂点を全て調べ済みの場合の処理
        u = stack.pop(-1)
        
        if len(stack) != 0:
            # 上流側の頂点に伝える
            upper = stack[-1]
            count[upper] += count[u]
            dist_sum[upper] += dist_sum[u] + count[u]   # 後ろの項は部分木内に含まれる各頂点へのupper⇔uの分

ans[1] = dist_sum[1]

# 深さ優先探索で隣接する頂点順に、部分木や計算済みの結果を利用して、他の根の場合での値を求める
# 頂点iが訪問済みかはansで代用
researched = [0 for _ in range(N+2)]        # 頂点iから移動できる頂点の何番目まで調べたか

stack = list()                              # 探索ルートを記録するstack
stack.append(1)

while len(stack) != 0:
    u = stack[-1]

    for k in range(researched[u],len(edge[u])):
        v = edge[u][k]      # uから辺がある頂点
        researched[u] += 1

        if ans[v] == -1:
            ans[v] = ans[u] - count[v] + (N - count[v])     # u→vしたので、v部分木内の頂点全て(count[v])1減。残りの分(N - count[v])は1増。cf2参照。
            stack.append(v)
            break   # 深さ優先なので、深い方に遷移出来たら抜ける
    else:
        # uから辺を張っている頂点を全て調べ済みの場合の処理
        u = stack.pop(-1)


# 回答出力
for i in range(1,N+1):
    print(ans[i])

最小全域木を求めたい

最小全域木 : 全頂点を連結する辺のコストが最小となるもの

クラスカル法

#https://qiita.com/Kept1994/items/051594a52dee5f8a4d3f

class UnionFind():
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        #要素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が属するグループと要素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):
        #要素xが属するグループのサイズ（要素数）を返す
        return -self.parents[self.find(x)]

    def same(self, x, y):
        #要素x, yが同じグループに属するかどうかを返す
        return self.find(x) == self.find(y)

    def members(self, x):
        #要素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 __str__(self):
        return '\n'.join(f'{r}: {m}' for r, m in self.all_group_members().items())


# 1.全ての辺を重みの小さいものから順に選びとる。
# 2.その過程で閉路を作ってしまうものは捨てる。全ての辺に対してこれを行う。

#入力読み込み
V, E = map(int, input().split())    #V:頂点数,E:辺数

edges = list()

for _ in range(E):
    s, t, w = map(int, input().split())     #s,t:点 , w:コスト
    edges.append((w, s-1, t-1))             #union findが0開始なので下駄を抜く。重みの小さいものから順に選びとるために、(重み,点,点)で全辺リストに追加

edges.sort(reverse=False)    #重みの小さい順で並び替え

uf = UnionFind(V)
cost = 0

#全ての辺を重みの小さいものから順に選びとる。
for edge in edges:
    w, s, t = edge

    if not uf.same(s, t):
        #点が既に同グループに属していない = 閉路を作らない ならOK
        cost += w
        uf.union(s,t)

print(cost)

有向グラフ中のループの有無を判定したい

強連結成分分解

# cf: https://hkawabata.github.io/technical-note/note/Algorithm/graph/scc.html

# 強連結成分分解,SCC,Strongly Connected Component
# = 有向グラフにおいて、互いに行き来が可能な頂点の集合

from collections import defaultdict

N = int(input())                        # 頂点数
M = int(input())                        # 有向辺の本数     

# 使用する辞書を確保
tree_dict = defaultdict(set)            # tree_dict[i]:ある頂点iから行ける頂点
inv_tree_dict = defaultdict(set)        # inv_tree_dict[i]:ある頂点iに行ける頂点

for i in range(M):
    A,B = map(str, input().split())     # 有向グラフの元,先

    # 頂点番号は0～N-1で処理するので、1始まりの場合はここで-1すること
    
    tree_dict[A].add(B)
    inv_tree_dict[B].add(A)

# 重複辺を消すためにdictの中身はsetで処理していたが、この後の処理のためにlistに変換する
tree_dict = [list(tree_dict[i]) for i in range(N)]
inv_tree_dict = [list(inv_tree_dict[i]) for i in range(N)]

# 2. SCCを求める
# 2.1 Gの頂点ラベリング
is_visited = [False for i in range(N)]    # その頂点が訪問済みか
researched = [0 for i in range(N)]        # その頂点の隣接点の何番目まで調べたか
label = list()                            # ラベル番号 index:ラベル番号,value:頂点番号

# 全点を対象に深さ優先探索
for i in range(0,N+1):

    if is_visited[i] == False:
        # 開始点の処理
        stack = list()
        stack.append(i)
        is_visited[i] = True
        
        while len(stack)!=0:
            #print(stack)
            node = stack[-1]
            connected_nodes = tree_dict[node]

            for j in range(researched[node],len(connected_nodes)):
                connected_node = connected_nodes[j]
                researched[node] += 1
                
                if is_visited[connected_node] == False:
                    is_visited[connected_node] = True
                    stack.append(connected_node)
                    break

            else:
                #行き止まり or 行先全部探索完了 ならラベリング
                label.append(stack.pop(-1))

# 2.2 逆グラフGTの探索
is_visited = [False for i in range(N+1)]    # その頂点が訪問済みか
researched = [0 for i in range(N+1)]        # その頂点の隣接点の何番目まで調べたか
SCC_list = list()

# ラベル番号最大(labelリストの最後尾かつ未訪問)のものから逆グラフGTをグラフ探索
while len(label)!=0:
    i = label.pop(-1)
    
    if is_visited[i] == False:
        SCC = set()
        
        stack = list()
        stack.append(i)
        
        is_visited[i] = True
        SCC.add(i)
        
        while len(stack)!=0:
            
            node = stack[-1]
            connected_nodes = inv_tree_dict[node]   # 探索対象は逆グラフGT
            
            for j in range(researched[node],len(connected_nodes)):
                connected_node = connected_nodes[j]
                researched[node] += 1
                if is_visited[connected_node] == False:
                    is_visited[connected_node] = True
                    stack.append(connected_node)
                    break
                    
            else:
                # 行き止まり or 行先全部探索完了 で SCCに加え(このタイミングでなくてもOK)、stackからpop
                SCC.add(stack.pop(-1))
        
        SCC_list.append(SCC)


# SCC_listにSCCがsetの形で入っているので利用して問題を解く

トポロジカルソートを試みて出来なければ閉路あり、出来れば閉路無しという判断も可能

LCA(最小共通祖先)を求めたい

e.g. ABC014-D
木構造で全点間最短距離を求めることが可能。

幅優先探索+ダブリングする

# https://algo-logic.info/lca/
# https://algo-logic.info/doubling/

from collections import deque

class ManagedTree():
    '''
    木構造に対して
    ・2頂点間の距離を求める
    ・ダブリングにより木の共通祖先(LCA)を求める
    ことを可能にする
    '''
    def __init__(self, root_node , edge):
        n = len(edge)                                       # n         : 木に含まれる頂点数
        self.dist = [-1] * n                                # dist[i]   : i⇔根の距離
        self.upstream_node = [list() for i in range(n+1)]   # upstream_node[i][j] : 頂点i の 2**j 上流側の頂点

        # 1. 幅優先探索で根からの距離と1段上流側の頂点を求めておく
        search_deq = deque([root_node])                     # 探索候補一覧
        self.dist[root_node] = 0
        self.upstream_node[root_node].append(root_node)

        while len(search_deq) != 0:
            i = search_deq.popleft()

            for j in edge[i]:
                if self.dist[j] == -1:
                    self.dist[j] = self.dist[i] + 1         # 全マス間距離1の場合。問題によって変えること
                    self.upstream_node[j].append(i)
                    search_deq.append(j)

        # 2. ダブリングを用いて 2**k 段上流側の頂点を求める
        max_dist = max(self.dist)

        j = 0

        while 1:
            for i in range(1,N+1):
                via_node = self.upstream_node[i][j]               # iから 2**j 段上流側の頂点
                reached_node = self.upstream_node[via_node][j]    # iから 2**j 上流側の頂点からさらに　2**j 段上流側の頂点 = iから 2**j + 2**j = 2 * 2**j = 2**(j+1) 段上流側の頂点
                self.upstream_node[i].append(reached_node)

            if 2**j >= max_dist:
                self.bit_len = j + 1    # 根⇔最深部の頂点の距離を2進数で表す際に必要なビット長
                break

            j += 1

    def get_upstream_node(self, n, k):
        # 頂点n の k段上流側の頂点 を求める
        if k == 0:
            return n
        else:
            for j in range(self.bit_len):
                if k & (1<<j):
                    n = self.upstream_node[n][j]
            return n

    def get_LCA(self, a, b):
        # 頂点a,b の LCA を求める
        # 1. 深い方の頂点を変更して、根からの距離を同じにする
        if self.dist[a] == self.dist[b]:
            d = 0
        else:
            if self.dist[a] > self.dist[b]:
                pass
            elif self.dist[a] < self.dist[b]:
                # 深い方の頂点をaとして扱う。
                b,a = a,b

            d = self.dist[a] - self.dist[b]

            for j in range(self.bit_len):
                if d & (1<<j):
                    a = self.upstream_node[a][j]

        if a == b:
            return a
        else:
            # 2. 二分探索を用いて、a,bの頂点から何段上がれば共通の頂点(LCA)となるかを求める
            maxi , mini = self.dist[a] , 0

            while 1:
                if maxi - mini <= 4:
                    for i in range(mini,maxi+1):
                        if self.get_upstream_node(a,i) == self.get_upstream_node(b,i):
                            return self.get_upstream_node(a,i)
                else:
                    mid = (maxi+mini) // 2        #int((max+min)//2)だと浮動小数点を介して、誤差になるので注意

                    if self.get_upstream_node(a,mid) == self.get_upstream_node(b,mid):
                        # mid 段上流側の頂点は一致したので、LCAはそれより下流にある
                        maxi = mid
                    else:
                        # mid 段上流側の頂点が不一致したので、LCAはそれより上流にある
                        mini = mid

    def get_dist(self, a, b):
        # 頂点a,b の 距離 を求める
        LCA = self.get_LCA(a, b)
        return self.dist[a] - self.dist[LCA] + self.dist[b] - self.dist[LCA]

N = int(input())
edge = [list() for _ in range(N+1)]

for i in range(N-1):
    x,y = map(int, input().split())
    edge[x].append(y)
    edge[y].append(x)

# 木を処理
root_node = 1
tree = ManagedTree(root_node,edge)

# クエリを処理する
Q = int(input())

for i in range(Q):
    a,b = map(int, input().split())
    print(tree.get_dist(a,b) + 1)

有向辺一覧を与えられたときに、条件を満たすグラフを求めたい

e.g. ABC291-E , ABC223-D

トポロジカルソートする。

# 辞書順最小のトポロジカルソートを求めるver
import heapq

# 1. 入力受付
N,M = map(int, input().split())

#制約をinput
restriction = [set() for i in range(N+1)]       # 自分がfromとなる有向辺 key:先に出てくる頂点 , value:後に出てくる頂点
inv_restriction = [set() for i in range(N+1)]   # 自分がtoとなる有向辺　 key:後に出てくる頂点 , value:先に出てくる頂点

for i in range(M):
    A,B = map(int, input().split())
    restriction[A].add(B)
    inv_restriction[B].add(A)

# 利用可能な頂点 = 制約のない頂点を求める
available = list()  # 利用可能な頂点一覧

for i in range(1,N+1):
    if len(inv_restriction[i]) == 0:
        available.append(i)


# 2. 回答導出
ans = list()                                # 辞書順最小となるトポロジカルソート
heapq.heapify(available)                    # 辞書順最小を取り出し & 途中の追加に対応させるためにheap化

while len(available) != 0:

    poped = heapq.heappop(available)        # 利用可能な頂点中で辞書順最小を取り出し
    ans.append(poped)                       # 回答最後尾に追加
    
    for v in restriction[poped]:            # 取り出した数字が掛けている制約全てに対して
        inv_restriction[v].remove(poped)    # 制約を消す
        
        if len(inv_restriction[v]) == 0:    # 頂点に対する制約が全てなくなっていれば
            heapq.heappush(available,v)     # → その頂点を利用可能に追加

# 3.回答出力
if len(ans) != N:
    print(-1)   # 閉路がある
else:
    print(' '.join(map(str,ans)))

～は～より先に現れる、とか、～は～より大きいとかの問題文のときも注意

部分木の値を用いて、木全体の値を求めたい

木DPする。
cf… https://algo-logic.info/tree-dp/

木DP e.g. ABC368_D

# Vの適当な頂点から始め深さ優先探索する
# 自分より下流側にVに入っている点が無ければ、その辺は削除できる

# 入力読み込み
N,K = map(int, input().split())

edge = [list() for _ in range(N+1)]
for i in range(N-1):
    A,B = map(int, input().split())
    edge[A].append(B)
    edge[B].append(A)

V = [int(e) for e in input().split()]



# 深さ優先探索の配列確保
is_visited = [False for _ in range(N+1)]    # 頂点iが訪問済みか
dp = [0 for _ in range(N+1)]                # 問題文を満たすような構成の木で、V[0]を根としたときの下流側の頂点数(自己を含む)
researched = [0 for _ in range(N+2)]        # 頂点iへ移動できる頂点の何番目まで調べたか
stack = list()                              # 探索ルートを記録するstack

# 深さ優先探索の初期化
S = V[0]
stack.append(S)                             # 開始点をstackに追加
is_visited[S] = True                        # 開始点を訪問済みにmark
V = set(V)                                  # 存在確認に使うので、高速になるようsetに変換

# 深さ優先探索の実施
while len(stack) != 0:
    u = stack[-1]                           # 上流側の点

    for k in range(researched[u],len(edge[u])):
        v = edge[u][k]                      # v : uから有向辺が張られている頂点(下流側の点)
        researched[u] += 1

        if is_visited[v] == False:
            # 当該頂点の下流側のdpは未計算なので、その方向に動く
            is_visited[v] = True
            stack.append(v)
            break                           # 深さ優先なので、深い方に遷移出来たら抜ける
    else:
        # uから有向辺を張っている頂点を全て調べ済みの場合の処理
        v = stack.pop(-1)                   # u : 葉(下流)側の頂点
        
        if len(stack) != 0:
            # まだ上流側に点がある → 探索継続。だがその前に木dpの値を根側の頂点に伝える
            u = stack[-1]                   # v: 根(上流)側の頂点

            # 自己がVに含まれている or 下流側に活かす点があるなら、自己を活かす必要があるので、その分を加える
            if v in V or dp[v] != 0:
                dp[v] += 1
            
            # 上流側の頂点に答えとなる頂点数を伝える
            dp[u] += dp[v]
        else:
            # 全探索完了したので、回答を出力
            print(dp[S] + 1)               # 開始点は必ずVに含まれているので、その分+1
            exit()

二部グラフか判定したい

二部グラフである ⇔ グラフの頂点を2色で塗り分けるとき、結ばれている点同士を異色に出来る塗り方が存在する
→ 適当な始点から幅優先探索で塗って、塗り方有無を確認する。e.g.ABC282-D
(※全体で連結でなくても、上記が成り立てば二部グラフであることに注意)

頂点同士で接続/切断したい

setを用いる。e.g. ABC302-E 「Isolation」
削除はO(1)、lenもO(1)なので十分高速。

集合を含むグラフを扱いたい

超頂点(集合としてまとめたい要素全てに対して辺を持つ頂点)を用いる。
e.g. ABC302-F

超頂点を使う

from collections import deque

N,M = map(int, input().split())

# 超頂点 … 各集合の点への辺を持つ点

is_visited = [False for _ in range(M+N+1)]      # 頂点(0～M)/超頂点(M+1～)を訪問済みか
ope_cnt = [-1 for _ in range(M+N+1)]            # 1から頂点iを同じ集合にするのに必要な最低操作回数
edge = [set() for _ in range(M+N+1)]            # 各頂点/超頂点の結びつき

search_deq = deque()                            #探索候補一覧

for i in range(N):
    A = int(input())
    S = [int(e) for e in input().split()]

    # 集合に含まれる頂点から超頂点へ辺を張る
    for j in S:
        edge[j].add(M+1+i)
    # 超頂点から含まれる頂点へ辺を張る
    edge[M+1+i] = S

# 幅優先探索：開始点の処理
search_deq.append(1)
ope_cnt[1] = -2     # 1と同集合の操作回数を0にするために少ない値から始める
is_visited[1] = True

# 幅優先探索：本処理
while len(search_deq) !=0 :
    i = search_deq.popleft()

    for j in edge[i]:
        if is_visited[j] == False:
            is_visited[j] = True
            ope_cnt[j] = ope_cnt[i] + 1
            search_deq.append(j)

if is_visited[M] == False:
    print(-1)
else:
    print(ope_cnt[M]//2)    # 超頂点を通る分の余分な増加を減らすために/2

確率・期待値を求めたい

確率・期待値の持つ性質を利用する
- 期待値 = 取りうる値とその値を取る確率の積の合計
- 事象AとBが独立のとき，Aが起きて，さらにBが起こる確率 = (Aが起きる確率)×(Bが起きる確率) …　積事象
- 事象AとBが同時に起きないとき、AかBのどちらが起こる確率 = (Aが起きる確率)+(Bが起きる確率) … 和事象
- 和の期待値 = 期待値の和
初期状態で確率が決まっているところ(勝確定、開始状態)から辿っていく
確率をmod～で求めてください → フェルマーの小定理から1/Nをmodしたものを求めて、/Nの代わりに利用する

～するまでの回数の期待値を求めたい

～するまでの回数の期待値by再帰(e.g.EDPC-J)

import pypyjit
pypyjit.set_param('max_unroll_recursion=0')
# ↑でPyPyの再帰が速くなるらしい https://qiita.com/shoji9x9/items/e7d19bd6f54e960f46be
# ★★★Pythonでの提出時はコメントアウトすること!!!!!★★★

import collections

N = int(input())
a = [int(e) for e in input().split()]

counted_a = collections.Counter(a)
u,v,w = counted_a[1],counted_a[2],counted_a[3]

# DPの配列確保：影響を及ぼすものを列挙し、その分を確保する。
# 今回は残り個数毎の皿数
dp = [[[-1 for _ in range(N+1)] for _ in range(N+1)] for _ in range(N+1)]
# dp[i][j][k]:皿に寿司が1,2,3個残っている皿数が(i,j,k)ときから(0,0,0)になるまでの回数の期待値

# DPの初期化：初期/終了状態で値が定まっているものがあれば、入れておく。
# 今回は、残り皿数が(0,0,0)のときは、当然回数0で(0,0,0)になるので、入れておく
dp[0][0][0] = 0

# 遷移を考える：ある状態 から/へ 変化できる状態を列挙する
'''
dp[i][j][k]から遷移できるパターンとその確率は
    dp[i-1][j][k]   : i / N
    dp[i+1][j-1][k] : j / N
    dp[i][j+1][k-1] : k / N
    dp[i][j][k]     : {N-(i+j+k)} / N
    (3,2個の皿から1個食べると、2,1の皿数が増えることに注意)
よって
    dp[i][j][k]
    = 1 +
      + dp[i-1][j][k] * (i / N)
      + dp[i+1][j-1][k] * (j / N)
      + dp[i][j+1][k-1] * (k / N)
      + dp[i][j][k]     * ({N-(i+j+k)} / N)
    (冒頭の1は次の状態に遷移するための分)

式を整理する。両辺にNを掛けて
    N * dp[i][j][k]
    = N
      + dp[i-1][j][k] * i
      + dp[i+1][j-1][k] * j
      + dp[i][j+1][k-1] * k
      + dp[i][j][k]     * {N-(i+j+k)}
右辺のdp[i][j][k]を移行し、
    dp[i][j][k] * (N - {N-(i+j+k)})
    = N
      + dp[i-1][j][k] * i
      + dp[i+1][j-1][k] * j
      + dp[i][j+1][k-1] * k
左辺を整理
    dp[i][j][k] * (i+j+k)
    = N
      + dp[i-1][j][k] * i
      + dp[i+1][j-1][k] * j
      + dp[i][j+1][k-1] * k
(i+j+k)で両辺を割る
    dp[i][j][k]
    = N / (i+j+k)
      + dp[i-1][j][k] * i / (i+j+k)
      + dp[i+1][j-1][k] * j / (i+j+k)
      + dp[i][j+1][k-1] * k / (i+j+k)
'''

# 個数の遷移がめんどうそうなので、メモ化再帰で解く
def rep(i,j,k):

    if dp[i][j][k] == -1:
        numerator = 0
        numerator += N
        if i != 0:
            numerator += rep(i-1,j,k) * i
        if j != 0:
            numerator += rep(i+1,j-1,k) * j
        if k != 0:
            numerator += rep(i,j+1,k-1) * k

        dp[i][j][k] = numerator / (i+j+k)
    
    return dp[i][j][k]

print(rep(u,v,w))

～する確率を求めたい

初期状態で確率が決まっているところ(勝確定、開始状態)から辿っていく

確率DP(後退解析) ABC298-E

# https://kyo-pro.hatenablog.com/entry/ABC298E
'''
    後退解析をする（先手の勝ち負けが確定している最終状態から，逆に考えていく）．
'''

mod = 998244353

N,A,B,P,Q = map(int, input().split())

# 1/P と 1/Q のmodしたものをフェルマーの小定理から求めておく
invP = pow(P, mod - 2, mod)
invQ = pow(Q, mod - 2, mod)

# DPする
# dpの配列確保：影響を及ぼすものを列挙し、その分を確保する。
# 今回は高橋君、青木君の位置:1～N、手番がどちらにあるか
dp = [[[0 for _ in range(2)] for _ in range(N+1)] for _ in range(N+1)]  # dp[i][j][k] : 高橋君がi、青木君がjにいて、k=0なら高橋君の手番、k=1なら青木君の手番、のときの高橋君の勝率

# DP配列の初期化：初期/終了状態で値が定まっているものがあれば、入れておく。
# 今回は既に高橋君がゴールにいれば勝利確定なのを利用する
for j in range(N+1):
    for k in range(2):
        dp[N][j][k] = 1

# 遷移を考える：ある状態 から/へ 変化できる状態を列挙する
# dp[i][j][0]からの遷移は出目1～P通りで、出目をdとすると遷移先での勝率は dp[min(N,i+d)][j][1] (制約よりオーバーでもゴール、手番が替わるのを考慮)
# それぞれの発生確率は1/P。
# ∴dp[i][j][0] = (1～Pまでの範囲で) dp[min(N,i+d)][j][1]/6 を足し合わせればOK

# 同様にdp[i][j][1]からの遷移は出目1～Q通り。それぞれの発生確率は1/Q。

for i in range(N-1,A-1,-1):
    for j in range(N-1,B-1,-1):
        for k in range(2):
            if k == 0:
                for d in range(1,P+1):
                    dp[i][j][k] += dp[min(N,i+d)][j][1]
                    dp[i][j][k] %= mod
                dp[i][j][k] *= invP
                dp[i][j][k] %= mod
            elif k == 1:
                for d in range(1,Q+1):
                    dp[i][j][k] += dp[i][min(N,j+d)][0]
                    dp[i][j][k] %= mod
                dp[i][j][k] *= invQ
                dp[i][j][k] %= mod

print(dp[A][B][0])

確率DP ABC323-F

MOD_BY = 998244353

N,X = map(int, input().split())
T = [int(e) for e in input().split()]

# 1/N のmodしたものをフェルマーの小定理から求めておく
inv_N = pow(N, MOD_BY - 2, MOD_BY)

# DPする
# dpの配列確保：影響を及ぼすものを列挙し、その分を確保する。
# 今回は時間のみ
dp = [0 for _ in range(X+1)]  # dp[i] : 時間iに次の曲が流れ始める確率

# DP配列の初期化：初期/終了状態で値が定まっているものがあれば、入れておく。
# 今回は時刻0から曲を流し始めることを利用する
dp[0] = 1

# 遷移を考える：ある状態 から/へ 変化できる状態を列挙する
# dp[i]からの遷移は流れる曲数N通り
# それぞれの発生確率は1/N。
# ∴dp[i + T[j]] = dp[i]/N を足し合わせればOK

ans = 0
for i in range(X+1):
    for j in range(N):
        if i + T[j] <= X:
            # 当該曲を流してもX+0.5を超えないので、流す
            dp[i + T[j]] += dp[i] * inv_N
            dp[i + T[j]] %= MOD_BY
        else:
            # 当該曲を流すとX+0.5を超える
            if j == 0:
                # 流れている曲が1なので、答えの確率に足す(和事象)
                ans += dp[i] * inv_N
                ans %= MOD_BY

print(ans)

~の期待値を求めたい

期待値DP ABC326-E

MOD_BY = 998244353

N = int(input())
A = [int(e) for e in input().split()]

# 1/N のmodしたものをフェルマーの小定理から求めておく
inv_N = pow(N, MOD_BY - 2, MOD_BY)

# DPする
# dpの配列確保：影響を及ぼすものを列挙し、その分を確保する。
# 今回はダイスの出目
dp = [0 for _ in range(N)]  # dp[i]:出目がyに到達する確率

# DP配列の初期化：初期/終了状態で値が定まっているものがあれば、入れておく。
# 今回はx=0を利用する

# 遷移を考える：ある状態 から/へ 変化できる状態を列挙する
# dp[i]からの遷移は流れる出目N通り
# それぞれの発生確率は1/N。
# ∴dp[j] = dp[i]/N (i<j)を足し合わせればOK

# 求めたいのは給料の期待値なので
# A[i] * dp[i] (給料額A[i] * iに到達できる確率)を足し合わせればOK

# が、dp[i]の値をi+1～Nまで配るようにするとTLEになる。
# ので貰う方で考える。
# dp[i] = dp[0]/N + dp[1]/N + ～ + dp[i-1]/N = (dp[0]～dp[i]の和)
# → 和を取りながら行うことによって間に合う

ans = 0     # 給料の期待値
dp_sum = 1  # i未満の出目の確率の和

for i in range(N):
    dp[i] = dp_sum * inv_N % MOD_BY
    dp_sum += dp[i]
    dp_sum %= MOD_BY

    ans += A[i] * dp[i]
    ans %= MOD_BY

print(ans%MOD_BY)

(条件を満たす)組み合わせを求めたい

貪欲法(値を最悪地で仮置きして、順番に最良値にしていく)

e.g. ABC362-C

(条件を満たす)最小/最大値を求めたい

幅優先探索・深さ優先探索
最小/最大全域木に帰着させてクラスカル法

貪欲法でシミュレーション(各グループから貪欲法で選びつつ、選ぶ個数を1個ずつずらす)

e.g.ABC312-F

N,M = map(int, input().split())

non_using_can_need_opener = list()                  # 選んでない缶切りが必要な缶
non_using_can_normal = list()                       # 選んでない缶切りが不要な缶
non_using_opener = list()                           # 選んでない缶切り

for i in range(N):
    T,X = map(int, input().split())

    if T == 0:
        non_using_can_normal.append(X)
    elif T == 1:
        non_using_can_need_opener.append(X)
    elif T == 2:
        non_using_opener.append(X)

non_using_can_normal.sort(reverse=False)
non_using_can_need_opener.sort(reverse=False)
non_using_opener.sort(reverse=False)

using_can_need_opener = list()                      # 選んでる缶切りが必要な缶
using_can_normal = list()                           # 選んでる缶切りが不要な缶 
using_opener = list()                               # 選んでる缶切り 

remain_open = 0                                     # 残りの開けられる缶の数

# 缶切り不要な方を優先して満足度の高い方からM個の缶を選ぶ
temp_ans = 0
for _ in range(min(M,len(non_using_can_normal))):
    highest_can = non_using_can_normal.pop(-1)
    temp_ans += highest_can
    using_can_normal.append(highest_can)

# まだM個に達していなければ、缶切りと缶切りが必要な缶も使う
while M > len(using_can_normal) + len(using_can_need_opener) + len(using_opener):
    if remain_open == 0:
        if len(non_using_opener) != 0:
            highest_opener = non_using_opener.pop(-1)
            using_opener.append(highest_opener)
            remain_open += highest_opener
        else:
            # 缶切りが必要な缶しか残ってないが、缶切りがないので開けられない
            break
    elif remain_open > 0:
        if len(non_using_can_need_opener) != 0:
            highest_can = non_using_can_need_opener.pop(-1)
            temp_ans += highest_can
            using_can_need_opener.append(highest_can)
            remain_open -= 1
        else:
            # 缶切りが必要な缶も残っていない
            break

ans = temp_ans

# 缶切り不要缶を抜きつつ、代わりに缶切り必要缶を入れていく。間で必要な缶切りも入れる
while len(using_can_normal) != 0:

    eject_can = using_can_normal.pop(-1)
    temp_ans -= eject_can
        
    if remain_open > 0:
        if len(non_using_can_need_opener) > 0:
            # 缶切りが必要な缶を開ける
            insert_can = non_using_can_need_opener.pop(-1)
            # using_can_need_opener.append(insert_can)
            remain_open -= 1
            temp_ans += insert_can
        else:
            # 缶切りはあるが、缶切りが必要な缶は残っていない
            break
    elif remain_open == 0:
        # 缶切りを補充する
        if len(non_using_opener) != 0:
            insert_opener = non_using_opener.pop(-1)
            # using_opener.append(insert_opener)
            remain_open += insert_opener
        else:
            # 缶切りが必要な缶は残っているかもしれないが、缶切りがないので開けられない
            break

    ans = max(ans,temp_ans)

print(ans)

DP

O(2**M)になりそうなときのDP e.g.tessokuA23

# N<=10なので取りうる状態は2**10 = 1024通りが最大
# 1024 * 100 = 102400 = 10**5 なので間に合いそう
# DPする

N,M = map(int, input().split())

# Aを2進数として扱い、それを数字に変換しつつ読込
A = [0 for i in range(M+1)]
for i in range(1,M+1):
    line = [int(e) for e in input().split()]
    num = 0
    for j in range(N):
        num += line[j] * pow(2,j)
    A[i] = num

# 初期化
dp = [[1000 for _ in range(2**N)] for _ in range(M+1)]    # dp[i][j]:i枚目までのクーポンを用いて状態jにするために最小クーポン枚数
dp[0][0] = 0

# DPする
for i in range(1,M+1):
    for j in range(2**N):
        dp[i][j] = min(dp[i][j],dp[i-1][j])                 # i番目のクーポンを不使用の場合
        dp[i][j|A[i]] = min(dp[i][j|A[i]],dp[i-1][j] + 1)   # i番目のクーポンを使用するの場合

if dp[M][2**N-1] == 1000:
    print(-1)
else:
    print(dp[M][2**N-1])

複数回使用不可条件があるときは、多重使用とならないように注意
- 更新前後の値に「～番目まで見たとき」の条件を入れる
- ループの順番を工夫する
- 更新後の値を別に取っておいて、最後にまとめて反映する

下か上に凸な関数の最小/最大値を求めたい

三分探索 e.g.ABC279-D

import math

def calc_ans(x):
    '''
    ある値での答えを求める
    '''
    return (A / math.sqrt(1+x)) + (B * x)

def search_ans(maxi,mini):
    '''
    mini～maxi間で判定用関数が最小値となるときの値を求める
    '''
    if maxi - mini <= 4:
        if calc_ans(mini) < calc_ans(mini+1):
            return mini
        else:
            for ans in range(mini,maxi+1):
                if calc_ans(ans) > calc_ans(ans+1) < calc_ans(ans+2):
                    return ans+1
            else:
                return maxi
    else:
        mid = (maxi+mini)//2        #int((max+min)//2)だと浮動小数点を介して、誤差になるので注意
        
        #判定用関数に投げて、条件により最大・最少を狭める
        if calc_ans(mid) < calc_ans(mid+1):
            # midより右に行くと値が増えるので最大値はmid以下
            maxi = mid
        else:
            # midより右に行くと値が減るので最小値はmid以上
            mini = mid
        
        return search_ans(maxi,mini)

A,B = map(int, input().split())
ans_cnt = search_ans(pow(10,16),0)
print(calc_ans(ans_cnt))

ある値の前後でYes/Noが分かれるときの最小/最大値を求めたい

二分探索 e.g.UTPC2020-A

#https://atcoder.jp/contests/utpc2020/tasks/utpc2020_a

def check_ans(T):
    '''
    ある値でOKか否かの判定関数。問題に合わせて変更する。
    '''
    pos = 0
    vitality = T
    
    for i in range(N):
        X,A = XA[i]

        vitality += X - pos
        if vitality > T:
            vitality = T
        
        vitality -= A
        if vitality < 0:
            return False

        pos = X    
    else:        
        return True

def binary_search(maxi,mini):
    '''
    mini～maxi間で判定用関数の結果がOKとなる最小値を求める
    '''
    
    if maxi - mini <= 4:
        for ans in range(mini,maxi+1):
            if check_ans(ans) == True:
                return ans
    else:
        mid = (maxi+mini)//2        #int((max+min)//2)だと浮動小数点を介して、誤差になるので注意
        
        #判定用関数に投げて、条件により最大・最少を狭める
        if check_ans(mid)==True:
            #midでOKだったので、OKとなる最小値はmidより小さい値である
            maxi = mid
        else:
            #midでNGだったので、OKとなる最小値はmidより大きい値である
            mini = mid
        
        return binary_search(maxi,mini)

N,L = map(int, input().split())

XA = list()
for i in range(N):
    X,A = map(int, input().split())
    XA.append((X,A))

print(binary_search(2*10**15,0))

小数で求めるとき　e.g. PAST10-G

# PyPyだと遅いのでループの回数に注意
from decimal import Decimal, getcontext
getcontext().prec = 12  # 必要な精度に変えること

a,b,c = map(Decimal, input().split())
maxi = Decimal('2.0')
mini = Decimal('1.0')

for i in range(50):
    x = (maxi+mini)/2
    
    if a * pow(x,5) + b * x + c >= 0:
        #midでOKだったので、OKとなる最小値はmidより小さい値である
        maxi = x
    else:
        #midでNGだったので、OKとなる最小値はmidより大きい値である
        mini = x

print(x)

区間を与えられて、最大に詰め込める個数を求めたい

区間スケジューリング問題 e.g.鉄則A39 https://atcoder.jp/contests/tessoku-book/tasks/math_and_algorithm_bn

# 終了時間が早いものを優先的に選ぶ(貪欲法)
import heapq

N = int(input())

movie_schedule = list()             # 上映スケジュールが(終了、開始)で入るリスト

for i in range(N):
    L,R = map(int, input().split())
    movie_schedule.append((R,L))    # 終了時間が早いものを選びたいので、前側のキーにして入れる

movie_schedule.sort(reverse=True)   # 終了時間が早いものが右に来るようにする

now_time = 0
ans = 0

while len(movie_schedule)!=0:
    R,L = movie_schedule.pop()

    if L >= now_time:
        # 開始時間が現在以降ならば当該の映画を見れる。終了時間が最速のものなので、最良の選び方
        now_time = R
        ans += 1

print(ans)

区間スケジューリング問題の応用(ABC325-D)

# 基本は区間スケジューリング問題なので、印刷範囲内から出てしまうのが速いものを優先すれば良い。
# e.g.鉄則A39 https://atcoder.jp/contests/tessoku-book/tasks/math_and_algorithm_bn
# 鉄則A39では終了時間までチャージ時間がかかるが、本問では1μsのみでOK
# → 印刷範囲内から出てしまうのが遅くても、印刷範囲内に先に入り印刷可能となる場合の考慮が必要 
# ∴ある時刻に印刷範囲内に入っているものを列挙→その中で印刷範囲内から出てしまうのが早いものを優先

import heapq

N = int(input())

products = list()                   # 商品の(印刷範囲内から出る時間、印刷範囲内に入る時間)を入れるリスト

for i in range(N):
    T,D = map(int, input().split())
    products.append((T,T+D))

products.sort(reverse=True)

now_time = 0
ans = 0
in_area_products = list()           # 印刷範囲内に入っている可能性のある商品が入るリスト
heapq.heapify(in_area_products)     

while 1:
    if len(in_area_products) == 0:
        # 印刷範囲内に入っているものがないので、時間を飛ばす
        if len(products) != 0:
            now_time = products[-1][0]

    # 現在時刻までに印刷範囲内に入っているものを、範囲内商品に加える
    while 1:
        if len(products) != 0 and products[-1][0] <= now_time:
            in_time,out_time, = products.pop()
            heapq.heappush(in_area_products,(out_time,in_time))
        else:
            break

    # 印刷範囲内商品で出て行ってしまうのが早いものを選ぶ
    while len(in_area_products) != 0:
        out_time,in_time = heapq.heappop(in_area_products)

        if in_time <= now_time <= out_time:
            now_time += 1
            ans += 1
            break

    if len(products) == 0 and len(in_area_products) == 0:
        break

print(ans)

(条件を満たす)n番目の～を求めたい

昇順に見て行くのが基本。数が多いときは、間の組み合わせをまとめて潰す。

(条件を満たす)n番目の値を求めたい

ダイクストラ法の応用(ABC297-E)

import heapq

N,K = map(int, input().split())
A = [int(e) for e in input().split()]

price_set = set()   # 既にカウント済みの支払える金額のset
hq = [A[i] for i in range(N)]
heapq.heapify(hq)
cnt = 0

while cnt < K:
    p = heapq.heappop(hq)   # 優先度付きキューを用いて最安値を取り出し
    
    if p not in price_set:
        # 同じ金額を支払う方法が複数存在する場合は1回だけ数えるようにする
        cnt += 1
        price_set.add(p)

        # 最安値に追加で1個買うのは最安値となる可能性があるので追加する
        for i in range(N):
            heapq.heappush(hq,p + A[i])

print(p)

条件を満たす組み合わせの個数を求めたい

DP

DP(ABC291-D)

MOD = 998244353

# 入力読み込み
N = int(input())
A_list = list()
B_list = list()

for i in range(N):
    A,B = map(int, input().split())
    A_list.append(A)
    B_list.append(B)

# DPテーブル準備
dp_front = [0 for i in range(N)]    # i番目のカードを表として使うとき、それまででいくつの並べ方があるか
dp_back = [0 for i in range(N)]     # i番目のカードを裏として使うとき、それまででいくつの並べ方があるか

# DP初期化:値があらかじめ分っている部分を埋める
dp_front[0] = 1
dp_back[0] = 1

# DP本処理:遷移式に従って、前の結果を用いて、後ろの答えを出していく
for i in range(1,N):
    if A_list[i] != A_list[i-1]:
        dp_front[i] += dp_front[i-1]
    if A_list[i] != B_list[i-1]:
        dp_front[i] += dp_back[i-1]
    
    if B_list[i] != A_list[i-1]:
        dp_back[i] += dp_front[i-1]
    if B_list[i] != B_list[i-1]:
        dp_back[i] += dp_back[i-1]
    
    dp_front[i] = dp_front[i] % MOD
    dp_back[i] = dp_back[i] % MOD

ans = dp_front[N-1] + dp_back[N-1]
print(ans % MOD)

Yes/Noで状態を表すことができ、順序が関係するならbit DPの可能性もあり
円環状の場合は、最初の値を決め打ちで考える(e.g.ABC307E)
1次元の場合は片方の端に着目、2次元の場合は隅に着目 (e.g.ABC311-E)

条件を満たす組み合わせでの最大値/最小値を求めたい

DP

DP(PAST16-H)

INF = pow(10,16)

# 入力読み込み
N,M = map(int, input().split())
A = [0] + [int(e) for e in input().split()]

# DPテーブル準備
dp_a = [[-INF for _ in range(M+1)] for _ in range(N+1)]     # i日目に宿題をして、それがj回目のときの最大の楽しさ
dp_b = [[-INF for _ in range(M+1)] for _ in range(N+1)]     # i日目に宿題をしないで、それまでにj回宿題しているときの最大の楽しさ 

# DP初期化:値があらかじめ判っている部分を埋める
# 今回は0日目部分を埋める
dp_a[0][0] = 0
dp_b[0][0] = 0

# DP本処理:遷移式に従って、前の結果を用いて、後ろの答えを出していく
for i in range(1,N+1):
    for j in range(M+1):
        if j >= 1:
            # j=0のとき、dp_b[i-1][j-1]でdp_b[i-1][-1]で末尾を読んでしまうので、if文で制約を掛ける
            dp_a[i][j] = dp_b[i-1][j-1]     # 2日連続で宿題しないので、dp_a[i-1][j-1]は呼ばない
        
        dp_b[i][j] = A[i] + max(dp_a[i-1][j],dp_b[i-1][j])
    
# 解答を求める
ans = max(dp_a[N][M],dp_b[N][M])
print(ans)

重複して使用してしまうパターンを避けるために、ループ2,3個目の数は最大→最小で回す!
[-i]で末尾からをi番目を参照させないようにif文で制約を掛けること!!

条件を満たす組み合わせの個数を求めたい

分割して考える

小さい条件を満たす場合を求める→小さい条件の組み合わせから全体の組み合わせ数を求める

ABC295-D

from collections import defaultdict

S = list(input())
N = len(S)

cnt = [0 for _ in range(10)]    # ある数字が何回出てきたか

flag_dict = defaultdict(int)    # 数字の出現回数の偶奇をまとめたもののdict
flag_dict[0] += 1

# 1.小さい条件を満たす場合を求める
for i in range(N):
    cnt[int(S[i])] += 1

    # 数字の出現回数の偶奇を求める
    flags = 0
    for j in range(10):
        flags += cnt[j]%2 * pow(10,j)
    
    flag_dict[flags] += 1

# 2.小さい条件を満たす場合の組み合わせから全体での答えを求める
ans = 0

for k,v in flag_dict.items():
    # 数字の出現回数の偶奇が一緒であることが嬉しい数列に出来る条件
    if v > 0:
        ans += v*(v-1)//2

print(ans)

DP

連続する部分文字列～の通り数を求めるときに使う。
部分文字列のうち、新たな文字が入る分、古い文字は忘れてよいことを利用する。

ABC359-D

MOD = 998244353

# 入力読み込み
N,K = map(int, input().split())
S = list(input())


# DPテーブル準備
dp = [[0 for _ in range(2**K)] for _ in range(N)]
# dp[i][j]:i+1文字目まで見たときに、末尾の文字列がjで、それまでに回文が存在しない通り数
# ただし、jはAを0,Bを1とした2進数表現を10進数に直したもの


# DP初期化:値があらかじめ判っている部分を埋める
# 今回は最初のK文字分を考える
for i in range(2**K):
    is_OK = True    # 最初のK文字でiに示す文字列を作れるか

    for j in range(K):
        if ((i >> (K-1-j)) & 1):
            # i の j bit目がONのときの処理
            if S[j] == "B" or S[j] == "?":
                pass
            else:
                is_OK = False
                break
        else:
            # i の j bit目がOFFのときの処理
            if S[j] == "A" or S[j] == "?":
                pass
            else:
                is_OK = False
                break
    
    if is_OK:
        dp[K-1][i] = 1


# K文字で回文になるものを調べて置く。
palindrome_set = set()

for i in range(2**K):
    is_palindrome = True    # iで示すK文字が回文であるか否か

    for j in range(K//2):   # 真ん中のbitは考慮しなくて良いので、K//2でOK
        if ((i >> (K-1-j)) & 1):
            if ((i >> (j)) & 1):
                pass
            else:
                is_palindrome = False
                break                
        else:
            if ((i >> (j)) & 1):
                is_palindrome = False
                break                
            else:
                pass
    
    if is_palindrome:
        # 最初のK文字で既に回文なら通り数を0にしておく
        dp[K-1][i] = 0
        palindrome_set.add(i)


# DP本処理:遷移式に従って、前の結果を用いて、後ろの答えを出していく
for i in range(K,N):
    for j in range(2**K):
        k = j & ~(1<<K-1)   # 左端の1文字分を削る
        k = (k<<1)          # ビットを左へシフト

        if S[i] == "B" or S[i] == "?":
            l = k           # l:S[i]によって新たに生成できる文字列。 
            l = l | (1<<0)  # 最右ビットを立てる

            if l not in palindrome_set:
                dp[i][l] += dp[i-1][j]
                dp[i][l] %= MOD
        
        if S[i] == "A" or S[i] == "?":
            l = k           # l:S[i]によって新たに生成できる文字列。
            # l = k | (1<<0)  # シフトした分は立っていないので、操作不要
            
            if l not in palindrome_set:
                dp[i][l] += dp[i-1][j]
                dp[i][l] %= MOD


# 解答を求める
ans = 0
for j in range(2**K):
    ans += dp[N-1][j]
    ans %= MOD

print(ans)

組み合わせを全探索/全列挙したい

多重ループ

半分全列挙

単純に全列挙するとTLEするものを半分ずつくらいに分けて、その中で全列挙してから組み合わせる

半分全列挙 e.g. 鉄則A14

from collections import defaultdict

N,K = map(int, input().split())

A = [int(e) for e in input().split()]
B = [int(e) for e in input().split()]
C = [int(e) for e in input().split()]
D = [int(e) for e in input().split()]

# N <= 1000 なので、そのまま全探索は 1000**4 = 10**12 でTLE
# A+B,C+Dの組み合わせに分ければ、2 * (1000**2) = 2 * 10**6 で間に合う

AB_sum = defaultdict(int)   # 組み合わせの有無なのでsetでも十分ではある
for i in range(N):
    for j in range(N):
        AB_sum[A[i] + B[j]] += 1

CD_sum = defaultdict(int)
for i in range(N):
    for j in range(N):
        CD_sum[C[i] + D[j]] += 1

for k,v in AB_sum.items():
    if K - k in CD_sum:
        print("Yes")
        exit()

print("No")

ON/OFFの組み合わせ

bit全探索

N = int(input())

# bit全探索
ans = 0
n = N   # n:bitが並ぶ個数

for i in range(2**n):
    part_ans = 0

    for j in range(n):
        if ((i >> j) & 1):
            # i の j bit目がONのときの処理
            pass
        else:
            # i の j bit目がOFFのときの処理
            pass
        
    ans = max(ans,part_ans)

print(ans)

N > 20 な場合には半分全列挙などを組み合わせる e.g. 典型90-#51

順序が関係するときはBit DP

状態がn個あるときの組み合わせの数を求めたい

nが2のときはbit全探索を参照

n進数全探索

X = int(input())
n = int(input())

# n進数全探索
for i in range(pow(n,X)):
    # iをn進数で表記する
    base_n_array = list()   # n進数表記のi(list形式)
    surplus = i
    for j in range(X-1,-1,-1):
        base_n_array.append(surplus//pow(n,j))
        surplus %= pow(n,j)

    print(base_n_array)

    # n進数表記の結果を用いて振り分け
    assignment = [list() for _ in range(n)]

    for j in range(X):
        assignment[base_n_array[j]].append(j)

    print(assignment)

n個から重複しないr個を選ぶ組み合わせ

n=26, r=13で $ {}_{n}C_r = 10,400,600 \fallingdotseq 10^7$。
問題例 : ABC018-D「バレンタインデー」 → 回答

【問題】
長さNの数列AからR個の数を選ぶとき、最大値はいくつになるか?

組み合わせの全探索

import itertools

# 入力読み込み
N,R = map(int, input().split())
A = [0] + [int(e) for e in input().split()]

# 組み合わせの全探索
ans = 0

for v in itertools.combinations([i for i in range(1,N+1)],R):
    # vが組み合わせ。n=6,r=3だと (1,2,3),(1,2,4),(1,2,5),...,(3,5,6),(4,5,6)
    part_ans = 0
    for i in v:
        # 組み合わせ内の各要素のアクセスして処理する
        part_ans += A[i]

    ans = max(ans,temp_ans)

print(ans)

単純に組み合わせの個数を求めるときは下記。(nCr)

組み合わせの個数

import math
 
def combinations_count(n, k):
    # math.comb(n,k)の代替。
    return math.factorial(n) // (math.factorial(n - k) * math.factorial(k))

n,r = map(int, input().split())
number_of_combination = math.comb(n,r)              # PyPy7.3.0だとRE
number_of_combination = combinations_count(n,r)     # PyPy7.3.0でもOK

n個から重複を許してr個を選ぶ組み合わせ(重複組み合わせ)を求めたい

n=14, r=14で $ {}_{n}H_r = 20,058,300 \fallingdotseq 2*10^7$。
問題例 : ABC268-D「Unique Username」 → 回答

【問題】
1～Nの数字が書かれたN枚のカードがある。1枚カードを引き、記録して、戻すことを5回繰り返す。
記録したカードがフルハウス(同じ数字のカードの3枚組と2枚組)となるパターンを列挙せよ。

重複組み合わせの全探索

import itertools
import collections

N = int(input())

for v in itertools.combinations_with_replacement([i for i in range(1,N+1)],5):
    # vが重複組み合わせ N=3,M=5で(1,1,1,1,1),(1,1,1,1,2),...,(1,1,3,3,3),...,(2,2,2,2,2),...,(3,3,3,3,3)
    counted_v = collections.Counter(v)
    
    if len(counted_v) == 2 and max(counted_v.values()) == 3:
        print(v)
        '''
        N = 3 で
        (1, 1, 1, 2, 2)
        (1, 1, 1, 3, 3)
        (1, 1, 2, 2, 2)
        (1, 1, 3, 3, 3)
        (2, 2, 2, 3, 3)
        (2, 2, 3, 3, 3)
        '''

n個からr個を選んで並べる組み合わせ(順列)を求めたい

n=11, r=11で $ {}_{n}C_r = 39,916,800 \fallingdotseq 4*10^7$。

【問題】
1~Nの数字が書かれたカードからR枚引く。引いたカードに書かれた数だけコマを進める。
スタートを1マス目としたとき、3の倍数が書かれたマスに止まれる最大の回数とその時のカードを求めろ。

組み合わせの全探索

import itertools

# 入力読み込み
N,R = map(int, input().split())

# 順列の全探索
ans = 0
ans_p = list()

for p in itertools.permutations([i for i in range(1,N+1)],R):
    # pが順列。n=6,r=3だと (1,2,3),(1,2,4),(1,2,5),...,(3,5,6),(4,5,6)
    part_ans = 0
    space_cnt = 1

    for i in p:
        # 組み合わせ内の各要素のアクセスして処理する
        space_cnt += i

        if space_cnt%3 == 0:
            part_ans += 1

    if ans < part_ans:
        ans = part_ans
        ans_p = p

print(ans,ans_p)

再帰を用いて全探索したい

再帰を用いた全探索 e.g.ABC326D

# 各行での並べ方はA,B,C,(.,.)なので最大 5*4*3 = 60通り
# かつ左にある文字の制約(例としてAとする)から、各行で試すのはそのうち
#   Aが先頭に来る場合 4*3*2 / 2 = 12
#  .Aが先頭に来る場合 3*2 = 6
# ..Aが先頭に来る場合 2 = 2
# で合計20通り
# 20**5 = 3200000 ≒ 3 * 10 ** 6通りなので全探索が間に合う

import itertools

N = int(input())
R = list(input())
C = list(input())

# まず60通りを全部列挙して、先頭文字毎に分ける
char_list = ["A","B","C"]
for _ in range(N-3):
    char_list.append(".")

pattern = {"A":list() , "B":list() , "C":list()}

for p in itertools.permutations(char_list,N):
    for i in range(N):
        if p[i] != ".":
            pattern[p[i]].append(p)
            break

# 各組み合わせを試す
A = [-1 for _ in range(N)]
col_head_cnt = [0 for _ in range(N)]
col_cnt = [{"A":0 , "B":0 , "C":0 , ".":0} for _ in range(N)]

def rep(n):
    # n+1行目を決める
    
    if n == N:
        # 各列の先頭文字が"."になっているパターンが残っているのでチェック
        for i in range(N):
            if col_head_cnt[i] == 0:
                return False
        else:
            return True    
    else:
        for p in pattern[R[n]]:
            # 列の最上文字条件・1文字のみ条件に抵触しないかチェック
            for i in range(N):
                if p[i] != ".":
                    if col_head_cnt[i] == 0 and C[i] != p[i]:
                        break
                    if col_cnt[i][p[i]] != 0:
                        break
            else:
                # 条件に抵触しなかったので、当該pで行を埋める。
                A[n] = p

                # 埋めた分、各列のカウントを増加する
                for i in range(N):
                    if C[i] == p[i]:
                        col_head_cnt[i] += 1
                    col_cnt[i][p[i]] += 1

                # 次行に進む。最後までOKならTrueを返す
                if rep(n+1):
                    return True

                # ダメだったので行を消して、その分の各列のカウントを戻す
                for i in range(N):
                    if C[i] == p[i]:
                        col_head_cnt[i] -= 1
                    col_cnt[i][p[i]] -= 1
    
    return False

if rep(0):
    print("Yes")
    for i in range(N):
        print(''.join(map(str,A[i])))
else:
    print("No")

二人ゲームの勝者を求めたい

再帰関数に落とし込んで、最良手を選ぶようにする

再帰関数への落とし込み&最良手を選ぶ e.g.鉄則A3

# 再帰を拡張しておく
import sys
sys.setrecursionlimit(1000000)

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

result = [[-1 for _ in range(2)] for _ in range(N+1)]   # result[i][j]:残り石個数がi個でターン数%2がjのときの敗者

def ref_func(remain,turn):
    if result[remain][turn%2] != -1:
        # 既に勝敗を求め済みならその値を返す
        return result[remain][turn%2]
    else:
        if remain < A and remain < B:
            # もう石を取れないので、今の手番側が敗者となる
            result[remain][turn%2] = turn % 2
        elif remain >= A and remain < B:
            # 石を取る方法は1つしかないので自動的に進む
            result[remain][turn%2] = ref_func(remain-A,turn+1)
        else:
            if turn%2 == 0:
                # 敗者が自分側(=0)にならないように、最大になる手を選ぶ
                result[remain][turn%2] = max(ref_func(remain-A,turn+1),ref_func(remain-B,turn+1))
            elif turn%2 == 1:
                # 敗者が自分側(=1)にならないように、最小になる手を選ぶ
                result[remain][turn%2] = min(ref_func(remain-A,turn+1),ref_func(remain-B,turn+1))

        return result[remain][turn%2]
    
ans = ref_func(N,0)

if ans == 0:
    print("Second")
else:
    print("First")

nimに落とし込む

nim : "N個の石山があり、交互に山から石を任意個取っていく。先に取れなくなったほうが負け"なゲーム
→ 各石山の個数をx[i]個とすると、勝敗は各山の個数を全てxorした値が=0なら先手は負け、!=0なら勝ち

# e.g. ABC368_F
'''
cf. 
https://yukicoder.me/problems/18
https://blog.hamayanhamayan.com/entry/2017/02/27/025050
https://note.com/kota_takahashi_/n/n54e837b3ec9d
https://algo-logic.info/combinatorial-games/#
'''
# 石を1個取る=Aに含まれる素因数1個で割る に換算することでnimに帰着できる。

from collections import defaultdict

def factorize(n):
    b = 2
    fct = defaultdict(int)
    while b * b <= n:
        while n % b == 0:
            n //= b
            fct[b] += 1
        b = b + 1
        
    if n > 1:
        fct[n] += 1
    
    return fct

N = int(input())
A = [int(e) for e in input().split()]

nim_A = list()
for i in range(N):
    factorized = factorize(A[i])
    nim_A.append(sum([i for i in factorized.values()]))

ans = nim_A[0]
for i in range(1,N):
    ans = ans^nim_A[i]

if ans == 0:
    print("Bruno")
else:
    print("Anna")

規則性を見つける

e.g. YUHA presents C88 謎解き×競技プログラミング『ある勇者の物語』B - ハヌマーンの試練

'''
残り枚数が3枚以下なら今手番がある方が勝ち
残り枚数が4枚なら今手番がある方が負け
残り枚数が5枚なら今手番がある方が勝ち 5→4→3→0
残り枚数が6枚なら今手番がある方が勝ち 6→4→3→0
残り枚数が7枚なら今手番がある方が勝ち 7→4→3→0
残り枚数が8枚なら今手番がある方が負け 8→(7,6,5)
残り枚数が9枚なら今手番がある方が勝ち 9→8→(7,6,5)
…
残り枚数が3枚以下or4の倍数でないなら今手番がある方が勝ち
'''

N = int(input())

if N <= 3 or N%4 != 0:
    print("SEN")
else:
    print("GO")

解法を思いつきたい

問題文を言い換えてみる

頂点iから最も近い～な点までの距離はd ⇔ 頂点iから距離d未満の範囲内に~な点は存在しない
+3,+5,+7する ⇔ -2,0,+2して、あとで操作回数x5を加える

変化量・不変量に着目する

数列内の数値を2つ選んでx,-xをそれぞれ加える → 数列全体の和は不変

細かく分けて考える

逆順に考える

後ろの方から考えると、うまくいく場合がある
e.g. ABC333-E

TLEを解消したい

異なる解法を試す
隠れた制約を使う
→　$ 0 < a < b < c $ かつ $a^2 * b * c^2 <=N　$ なら $a+1 < b , a+2 < c $なので $b≒a,c≒a$と考えて、と $a^5 <= N$ ... e.g.ABC300-E
→ 幅優先探索で探索範囲を限定する
多重ループをΣの性質を使って分解する
i<j<k 制約がついているとき、後ろから処理する。　e.g.ABC308E
- 同値が許容できない場合は貯めておく
- 残り分のループをセグメント木やbisectなどで処理する　e.g.ABC309F
制約が二重にある e.g. 平面上である座標より左上での最大値を求めたいABC369F
- sortして片方の制約を満たすようにしつつ、もう片方の制約をseg木で探す

PyPyを使う

再帰は遅いので下記を冒頭に入れる

import pypyjit
pypyjit.set_param('max_unroll_recursion=0')
# Pythonでの提出時はコメントアウトすること!!

その他

衝突して向きを変える問題を解きたい

人が入れ替わるだけで、そのまま移動するものとして扱ってOK
e.g. 鉄則A43

いろいろなDP

桁DP

上or下の桁から決めていくDP。
https://twitter.com/june19312/status/1646685054800916480?s=20

桁DP e.g.ABC154E

import math

def combinations_count(n, k):
    # math.comb(n,k)の代替。
    return math.factorial(n) // (math.factorial(n - k) * math.factorial(k))

N = int(input())
K = int(input())

# Nの桁数に達するまでの処理
ans = 0
digit = 1

while 1:
    digit_max = pow(10,digit) - 1
    digit_min = pow(10,digit-1)

    if digit_max > N:
        break
    elif digit - 1 >= K - 1:
        digit_all = digit_max - digit_min + 1
        # ans += pow(9,K) * math.comb(digit-1,K-1)    # 0でない箇所への数字の入れ方の通り * 頭以外の場所からK-1個、0でない数字にする箇所を選ぶ通り。PyPy7.3.0だとREなので下記で書く。
        ans += pow(9,K) * combinations_count(digit-1,K-1)
    
    digit += 1

# N = 99...99 の時は継続処理の必要なし
if digit_min - 1 == N:
    print(ans)
    exit()

# digit_min ～ Nまでの分を桁DPで求める。
# Nを各桁のリスト形式に変換
N = list(str(N))
N = [int(N[i]) for i in range(len(N))]

# 桁DPの配列宣言
same = [[0 for _ in range(K+1)] for _ in range(len(N))]     # i桁目まで見たときにNと一致する通り数の中で0でない数をj個含むもの
under = [[0 for _ in range(K+1)] for _ in range(len(N))]    # i桁目まで見たときにN未満になる通り数の中で0でない数をj個含むもの

# 頭の桁の処理
under[0][1] = N[0] - 1  # 0を除くNの頭の桁未満の分
same[0][1] = 1          # Nの頭の桁と一致なので1通り

# それ以外の桁の処理
for i in range(1,len(N)):
    if N[i] == 0:
        for j in range(K+1):
            same[i][j] += same[i-1][j]                      # i桁目を一致できるのは"0" 1通り。i桁目を一致させるのには0を使うので、0以外の使用数はそのままとなる。
    else:
        for j in range(1,K+1):
            same[i][j] += same[i-1][j-1]                    # i桁目を一致できるのはNのi桁目1通り。i桁目を一致させるのに0以外の値を使う分を考慮する

    for j in range(K+1):
        if j-1 >= 0:
            under[i][j] += under[i-1][j-1] * 9              # i桁目より左側で下回っていれば、自由に数字(1～9)を選べる。0以外の値を使っているので、その分を考慮する

        under[i][j] += under[i-1][j]                        # i桁目より左側で下回っていて、i桁目に0を使うパターン。0を使っているので、0以外の使用数はそのままとなる。

        if N[i] != 0:
            if j-1 >= 0:
                under[i][j] += same[i-1][j-1] * (N[i]-1)    # i-1桁目より左側が同じとき、i桁目まで見たときにN以下にするパターン。Nのi桁目を下回る数が通り数。0以外の値を使用することを考慮する。
            under[i][j] += same[i-1][j]                     # i-1桁目より左側が同じとき、i桁目を0とし、N以下にするパターン。0以外の使用数はそのままとなる。

ans += same[len(N)-1][K] + under[len(N)-1][K]
print(ans)

二次元DP

sortして片方の制約を満たすようにしつつ、もう片方の制約をseg木で探す
seg木に値だけではなく、indexも記録しておくことで、順番も探せる

行列累乗 e.g.ABC369F

# コインの位置を昇順に並び替える
# 左から順に処理していけば、行は必ず自分より若いものになる
# かつseg_treeを用いることにより、列が自分より若い中での最大値も求まる

INF = 10**6

# https://qiita.com/R_olldIce/items/32cbf5bc3ffb2f84a898
class SegTree:
    def __init__(self, op, e, n, v=None):
        self._n = n
        self._op = op
        self._e = e
        self._log = (n - 1).bit_length()
        self._size = 1 << self._log
        self._d = [self._e()] * (self._size << 1)
        if v is not None:
            for i in range(self._n):
                self._d[self._size + i] = v[i]
            for i in range(self._size - 1, 0, -1):
                self._d[i] = self._op(self._d[i << 1], self._d[i << 1 | 1])

    def set(self, p, x):
        # 更新クエリ: a[p] を x に更新
        p += self._size
        self._d[p] = x
        while p:
            self._d[p >> 1] = self._op(self._d[p], self._d[p ^ 1])
            p >>= 1

    def get(self, p):
        # p が与えられたとき、a[p] を取得する
        return self._d[p + self._size]

    def prod(self, l, r):
        # 取得クエリ: 区間 [l,r) の総積を取得 (※ここでの積は"掛け算"でなく"演算結果"くらいの意味)
        # 【注意!】: ")" なので区間の終わりの数字ちょうどは含まれない! 必要に応じて+1すること!
        sml, smr = self._e(), self._e()
        l += self._size
        r += self._size
        while l < r:
            if l & 1:
                sml = self._op(sml, self._d[l])
                l += 1
            if r & 1:
                r -= 1
                smr = self._op(self._d[r], smr)
            l >>= 1
            r >>= 1
        return self._op(sml, smr)

    def all_prod(self):
        # 全要素の総積を取得 (※ここでの積は"掛け算"でなく"演算結果"くらいの意味)
        return self._d[1]

def op(x, y):
    # 演算
    xv,xi = x
    yv,yi = y
    if xv > yv:
        return (xv,xi)
    else:
        return (yv,yi)

def e():
    # 単位元。任意の要素に対して演算をしても影響を与えないような値。
    # opが加算なら0,掛け算なら1,minならINF,maxなら-INF,xorなら0,andなら1,orなら0
    # 今回は値と前の座標番号を記録する
    return (-1,-1)  


# 入力受付
H,W,N = map(int, input().split())

coins = list()
for i in range(N):
    R,C = map(int, input().split())
    R -= 1      # 0始まりにするためにshift
    C -= 1      # 同上
    coins.append((R,C))


# セグメントツリーを用いて、若い列から処理することにより、解答を求める
init_list = [(0,-1)] *(W+2)                         # 初期化用リスト
seg = SegTree(op, e, len(init_list) , init_list)    # 左から演算,単位元,初期値リストの長さ,初期値リスト。今まで見たcoinの座標中、最大コイン取得数と、indexを記録

prev_node = [-1 for _ in range(N)]    # ルートを記録する

coins.sort(reverse=False)
for i in range(N):
    R,C = coins[i]
    val,index = seg.prod(0,C+1)         # 自分より左上の中で最もコインを取れた枚数と、そのindexを求める。(sortしているので、行側は自動的に成立している)
    seg.set(C , (val+1,i))              # 自分の分+1して、seg treeに記録
    prev_node[i] = index


# 解答を形式に合わせる
val,index = seg.prod(0,W+2)
print(val)

# ルートを求める。最後のコインから順々に辿る
route = list()
m,n = H-1,W-1                           # 0始まりにしているので、ゴール地点もシフトする

while 1:
    u,v = coins[index]      
    route += ["D"] *(m-u)
    route += ["R"] *(n-v)

    index = prev_node[index]            # 1枚前のコインのindexを得る
    m,n = u,v
    if index == -1:
        # 前のコインがない → 移動前はスタート地点なので、当該コインからスタート地点までをたどる
        u,v = 0,0
        route += ["D"] *(m-u)
        route += ["R"] *(n-v)
        break

print(''.join(map(str,route[::-1])))

操作/変化を繰り返した結果を求めたい

再帰する

行列累乗する

DPの漸化式を行列の形式にし、遷移の回数分行列を累乗し、初期値と掛けて答えを求める。

行列累乗 e.g.DPまとめコンテスト-R　"Walk"

# cf.https://www.npca.jp/magazine/2021/Sugsugar

from collections import deque

MOD_BY = pow(10,9) + 7

def get_matrix_product(A,B):
    '''
    行列Aと行列Bの積を求める
    # cf : https://w3e.kanazawa-it.ac.jp/math/category/gyouretu/senkeidaisu/henkan-tex.cgi?target=/math/category/gyouretu/senkeidaisu/gyouretu-no-seki.html
    '''
    l = len(A)        # 行列Aの行数
    m = len(A[0])     # 行列Aの列数 (= 行列Bの行数)
    n = len(B[0])     # 行列Bの列数
    
    if m != len(B):
        # 行列Aの列数と行列Bの行数が不一致のため、行列積を計算できない
        return -1
    else:
        C = [[0 for _ in range(n)] for _ in range(l)]

        for i in range(l):
            for j in range(n):
                for k in range(m):
                    C[i][j] += A[i][k] * B[k][j]
                    C[i][j] %= MOD_BY
    
        return C

def get_power_matrix(A,n):
    '''
    A**nの行列を求める
    '''
    if len(A) != len(A[0]):
        # 正方行列でなければ、A**Nは計算不可
        return 0
    
    size = len(A)   # Aの行数 (= Aの列数)

    if n & (1<<0):
        # A ** N は A**1を含むので、Aで仮置き
        ans = [[A[i][j] for j in range(size)] for i in range(size)]
    else:
        # A ** N は A**1を含まないので、単位行列で仮置き
        ans = [[0 for _ in range(size)] for _ in range(size)]
        for i in range(size):
            ans[i][i] = 1

    # A**2,A**4,A**8 ... を求めつつ,nの当該bitが立っていたらansに掛け合わせる
    powered_A = A

    for i in range(1,64):
        powered_A = get_matrix_product(powered_A,powered_A)     # A ** (i+1)
    
        if n & (1<<i):
            # A ** N は A**(i+1)を含むので、ansに掛け合わせる
            ans = get_matrix_product(powered_A,ans)
    
    return ans

N,K = map(int, input().split())

# 長さがK = 遷移回数がK回。Kが大きいので行列累乗する

# 遷移を示す行列aを定める。今回は入力で与えられているものを利用する。
a = [[int(e) for e in input().split()] for _ in range(N)]
a = list(zip(*a))                       # 与えられたaをdp[i] = a・dp[i-1] になるように転置する

# (必要であれば)答えを求めるのに必要な累乗となるように、Kの値を調整。

# 行列累乗本体
a_K = get_power_matrix(a,K)             # a ** K
dp_0 = [[1] for _ in range(N)]          # 遷移式の初期値。今回は初回に各頂点にいる通り数
dp_K = get_matrix_product(a_K,dp_0)     # dp[K] = a・dp[K-1]　= a・a・dp[K-2] = … = a**K・dp[0]

# 設問から答えは合計値なので求めて出力する
ans = 0
for i in range(N):
    ans += dp_K[i][0]
    ans %= MOD_BY

print(ans)

ダブリングする

ダブリング e.g.鉄則A57

N,Q = map(int, input().split())
A = [0] + [int(e) for e in input().split()]

XY = list()
max_Y = 0
for i in range(Q):
    X,Y = map(int, input().split())
    XY.append((X,Y))
    max_Y = max(max_Y,Y)

doubling = [list() for _ in range(N+1)]         # doubling[i][j]:今穴iにいるアリが2**j日後にどの穴にいるか
for i in range(1,N+1):
    doubling[i].append(A[i])                    # 2**0 = 1日後にいる穴はAで記されているので代入

j = 0
while 1:
    for i in range(1,N+1):
        via_hole = doubling[i][j]               # 2**j日後にいる穴
        reached_hole = doubling[via_hole][j]    # 2**(j+1)日後にいる穴
        doubling[i].append(reached_hole)
    
    if pow(2,j) >= max_Y:
        # 最大値を表せるようになったら終了
        bit_len = j+1    # 最も大きいYを表すのに必要なbit長
        break
    else:
        j += 1

for i in range(Q):
    X,Y = XY[i]

    # Y を 2**n の和で考える
    for j in range(bit_len):
        if Y & (1<<j):
            X = doubling[X][j]

    print(X)

ループする場合、ループ長で余りを取る

変化する箇所・しない箇所を考えて、する→しないの変化を考える

e.g. ARC024-B・・・解答