pres1dent.net
/

    アルゴリズム

    Python 競技プログラミング実戦リファレンス

    このページは Python で問題を解くときの「思い出すための索引」である。文法の入門ではなく、計算量と実装上の罠を含む実戦用メモとして使う。

    最小テンプレート

    import sys

input = sys.stdin.readline


def solve() -> None:
    n = int(input())
    a = list(map(int, input().split()))
    print(sum(a))


if __name__ == "__main__":
    solve()

    input を差し替える場合、返り値の末尾に改行が残る。ただし int()、map(int, ...)、split() に渡すなら通常は問題にならない。

    入力

    一行の整数

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

    map は遅延イテレータである。一度消費すると空になるので、複数回参照するなら list にする。

    複数行の配列

    h, w = map(int, input().split())
grid = [input().rstrip() for _ in range(h)]

n = int(input())
points = [tuple(map(int, input().split())) for _ in range(n)]

    文字列グリッドでは rstrip() を使う。strip() は左右の空白も消すため、空白自体がデータの場合に意味が変わる。

    入力全体を一度に読む

    整数しか現れず、行の境界が不要なら一括読み込みが簡潔で速い。

    it = iter(map(int, sys.stdin.buffer.read().split()))
n = next(it)
values = [next(it) for _ in range(n)]

    固定個数ずつまとめる例:

    it = iter(map(int, sys.stdin.buffer.read().split()))
n = next(it)
edges = [(next(it), next(it)) for _ in range(n - 1)]

    EOF まで読む

    for line in sys.stdin:
    if not line.strip():
        continue
    a, b = map(int, line.split())

    一括読み込みなら次の形でもよい。

    data = list(map(int, sys.stdin.buffer.read().split()))

    出力

    列と行列

    print(*values)
print("\n".join(map(str, values)))
print("\n".join("".join(row) for row in grid))

    大量の行をループで print するより、文字列へまとめて一度に出力する方が速いことが多い。

    Yes / No

    print("Yes" if condition else "No")

    表記は問題文に合わせる。YES と Yes は別の出力である。

    浮動小数点

    print(f"{answer:.15f}")

    許容誤差がある問題では十分な桁を出す。金額や組合せの厳密値に浮動小数点を使わない。

    リストと文字列

    初期化の罠

    # 正しい: 行ごとに異なるリスト
grid = [[0] * w for _ in range(h)]

# 誤り: 全行が同じリストを参照する
grid = [[0] * w] * h

    イミュータブルな値だけを並べる一次元リストでは [0] * n で問題ない。

    スライス

    a[l:r]        # 半開区間 [l, r)
a[::-1]       # 逆順のコピー
a[::2]        # 偶数添字
a[1::2]       # 奇数添字

    スライスは要素数に比例するコピーを作る。ループ内で大きなスライスを繰り返すと二次時間になりやすい。

    文字列の構築

    文字列への += を大量に繰り返すより、断片をリストへ集めて join する。

    parts = []
for x in values:
    parts.append(str(x))
answer = "".join(parts)

    collections

    deque: キューと両端操作

    from collections import deque

que = deque([start])
while que:
    v = que.popleft()
    for to in graph[v]:
        if dist[to] == -1:
            dist[to] = dist[v] + 1
            que.append(to)

    リストの pop(0) は後続要素をずらすので O(N)O(N)。deque.popleft() は O(1)O(1)。

    0-1 BFS では、コスト 0 の遷移を左、コスト 1 の遷移を右へ追加する。

    if cost == 0:
    que.appendleft(to)
else:
    que.append(to)

    Counter: 頻度表

    from collections import Counter

count = Counter(values)
most_common = count.most_common()

    存在しないキーの値は 0。辞書として走査するとキーを列挙する。

    for value, frequency in count.items():
    ...

    defaultdict: 隣接リストや分類

    from collections import defaultdict

groups = defaultdict(list)
for key, value in pairs:
    groups[key].append(value)

    読み取りだけのつもりで groups[key] を呼ぶと、キーが新規作成される。存在確認には key in groups を使う。

    ヒープ

    heapq は最小ヒープである。

    import heapq

heap = []
heapq.heappush(heap, (distance, vertex))
distance, vertex = heapq.heappop(heap)

    タプルは先頭から辞書順に比較される。ダイクストラ法 では (距離, 頂点) を積む。

    最大ヒープ

    値の符号を反転する方法は、古い Python を含めて使える。

    heapq.heappush(heap, -value)
maximum = -heapq.heappop(heap)

    上位 K 個だけ保持する

    heap = []
for x in values:
    if len(heap) < k:
        heapq.heappush(heap, x)
    elif x > heap[0]:
        heapq.heapreplace(heap, x)

    計算量は O(Nlog⁡K)O(N\log K)、追加領域は O(K)O(K)。

    ソート

    Python のソートは安定である。同じキーを持つ要素の元の順序が保たれる。

    items.sort()                              # 破壊的、返り値は None
ordered = sorted(items)                  # 新しいリスト
items.sort(key=lambda x: (x[1], -x[0]))  # 第2要素昇順、第1要素降順

    座標圧縮

    ordered = sorted(set(values))
rank = {value: i for i, value in enumerate(ordered)}
compressed = [rank[value] for value in values]

    辞書の構築まで O(Nlog⁡N)O(N\log N)、各変換は平均 O(1)O(1)。

    二分探索 bisect

    from bisect import bisect_left, bisect_right

left = bisect_left(a, x)    # x 以上となる最初の位置
right = bisect_right(a, x)  # x より大きくなる最初の位置
count = right - left        # x の個数

    a は昇順でなければならない。探索は O(log⁡N)O(\log N) だが、リスト中央への insort は挿入後の移動があるため O(N)O(N)。

    境界値そのものを探索する問題は 答えで二分探索 を参照する。

    集合と辞書

    set と dict の検索・追加・削除は平均 O(1)O(1)。順序付き集合ではない。

    seen = set()
seen.add(x)
seen.discard(x)  # 存在しなくても例外にならない

if x in seen:
    ...

    remove(x) は存在しない場合に KeyError、discard(x) は何もしない。

    集合演算

    a | b   # 和集合
a & b   # 積集合
a - b   # 差集合
a ^ b   # 対称差

    大きい集合に対する演算は集合サイズに比例する。ループ内で新しい集合を何度も作らない。

    整数と剰余

    Python の整数は任意精度なのでオーバーフローしないが、値が巨大になるほど演算コストは増える。

    q, r = divmod(a, m)
power = pow(a, exponent, mod)
inverse = pow(a, -1, mod)  # 逆元が存在する場合

    法が素数 pp で a≢0(modp)a \not\equiv 0 \pmod p なら、フェルマーの小定理から pow(a, p - 2, p) でも逆元を求められる。

    ビット操作

    bit = 1 << i
has_i = (mask >> i) & 1
mask |= 1 << i       # 追加
mask &= ~(1 << i)    # 削除
mask ^= 1 << i       # 反転
count = mask.bit_count()

    部分集合全探索:

    for mask in range(1 << n):
    selected = [i for i in range(n) if mask >> i & 1]

    ある集合 mask の部分集合を列挙する:

    sub = mask
while True:
    # sub を処理
    if sub == 0:
        break
    sub = (sub - 1) & mask

    グラフ

    隣接リスト

    graph = [[] for _ in range(n)]
for _ in range(m):
    u, v = map(int, input().split())
    u -= 1
    v -= 1
    graph[u].append(v)
    graph[v].append(u)

    木の親・深さ・走査順を反復 DFS で作る:

    parent = [-1] * n
depth = [0] * n
order = [0]
for v in order:
    for to in graph[v]:
        if to == parent[v]:
            continue
        parent[to] = v
        depth[to] = depth[v] + 1
        order.append(to)

    部分木の値を集約するときは reversed(order) で子から親へ戻る。詳しくは 木の走査。

    計算量と性能感覚

    制約を見たら、実装前に許される計算量を決める。

    NN の目安典型的に検討する計算量
    N≤20N \le 20O(2NN)O(2^N N)、半分全列挙
    N≤400N \le 400O(N3)O(N^3)
    N≤3×103N \le 3\times 10^3O(N2)O(N^2)
    N≤2×105N \le 2\times 10^5O(Nlog⁡N)O(N\log N)、O(N)O(N)
    N≥107N \ge 10^7小さい定数の O(N)O(N) または数式化

    これは保証ではない。Python では同じ計算量でも、辞書操作、オブジェクト生成、関数呼び出しの回数によって大きく変わる。

    よく効く改善

    1. 漸近計算量を下げる
    2. 同じ値の再計算を避ける
    3. 内側のループから辞書・関数呼び出しを減らす
    4. 大量入力をまとめて読む
    5. 不要なコピーと一時リストを作らない

    局所的な短縮より、まずアルゴリズムとデータ構造を見直す。区間集約には セグメント木、単調な境界には 答えで二分探索 のように、必要な操作から選ぶ。

    デバッグ用断片

    標準エラー出力

    print("debug:", value, file=sys.stderr)

    提出出力と混ざらないが、不要なデバッグコードは消す。

    小さいケースとの比較

    最適化した解法と愚直解をランダムな小入力で比較する。

    import random

for n in range(1, 9):
    for _ in range(1000):
        a = [random.randrange(10) for _ in range(n)]
        expected = brute_force(a)
        actual = solve_fast(a)
        assert actual == expected, (a, expected, actual)

    境界条件だけを人手で列挙するより、反例が存在する領域を広く探索できる。特に貪欲法、尺取り法、二分探索の境界確認に有効である。

    提出前チェック

    • 添字を 0-indexed に直したか
    • 無向辺を両方向へ追加したか
    • 空配列、単一要素、到達不能を扱えるか
    • INF は答えの上限より十分大きいか
    • 浮動小数点を厳密比較していないか
    • 再帰の深さは安全か
    • 計算量とメモリ量を制約から確認したか
    • サンプル以外の最小ケースを試したか