코딩테스트용 파이썬 알고리즘 치트시트: 복붙해서 바로 쓰는 핵심 스니펫 모음

코딩테스트나 PS(Problem Solving)를 하다 보면 매번 똑같은 코드를 검색하게 됩니다. "파이썬 이진 탐색", "파이썬 다익스트라 템플릿", "파이썬 빠른 입력"… 결국 검색 → 복붙 → 약간 수정의 무한 반복이죠.

이 글은 그 검색을 한 번에 끝내기 위한 레퍼런스입니다. 자주 쓰는 스니펫만 모았고, 전부 복사해서 바로 붙여넣을 수 있도록 정리했습니다. 북마크해두고 필요할 때마다 꺼내 쓰세요.

동작 환경: Python 3.8+ (표준 라이브러리만 사용, 별도 설치 불필요)

---

1. 빠른 입력 / 출력 (Fast I/O)

input()을 그냥 쓰면 입력이 많을 때 시간 초과(TLE)가 납니다. 반복 입력이 많은 문제라면 거의 항상 이걸로 시작하세요.

import sys
input = sys.stdin.readline   # 이후 input()이 빨라짐

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

readline()은 개행 문자(\n)까지 같이 읽으므로, 문자열을 그대로 쓸 땐 .rstrip()을 붙이세요.

s = input().rstrip()   # 끝의 개행 제거

출력이 많을 땐 한 번에 모아서 출력하는 게 훨씬 빠릅니다.

import sys
print = sys.stdout.write     # 주의: 문자열만, 개행 직접 추가

# 또는 결과를 모았다가 한 번에
result = []
for x in range(5):
    result.append(str(x))
sys.stdout.write("\n".join(result) + "\n")

---

2. 재귀 깊이 제한 풀기

파이썬 기본 재귀 한도는 1000입니다. DFS를 재귀로 짜면 깊은 그래프에서 RecursionError가 터집니다.

import sys
sys.setrecursionlimit(10**6)

너무 큰 값을 주면 세그멘테이션 폴트가 날 수 있습니다. 보통 10**6 정도면 충분합니다. 깊이가 매우 깊다면 재귀 대신 스택 기반 반복문으로 바꾸는 걸 권장합니다.

---

3. 이진 탐색 (bisect)

직접 구현할 필요 없이 표준 라이브러리 bisect를 쓰는 게 정석입니다.

import bisect

arr = [1, 3, 3, 5, 7]

bisect.bisect_left(arr, 3)    # 3이 들어갈 가장 왼쪽 인덱스 → 1
bisect.bisect_right(arr, 3)   # 3이 들어갈 가장 오른쪽 인덱스 → 3

# 특정 값의 개수 = 오른쪽 경계 - 왼쪽 경계
count = bisect.bisect_right(arr, 3) - bisect.bisect_left(arr, 3)  # → 2

# 정렬 상태를 유지하며 삽입
bisect.insort(arr, 4)         # arr이 정렬된 채로 4 삽입

직접 구현 버전이 필요할 때 (예: "특정 조건을 만족하는 최솟값" 매개변수 탐색):

def binary_search(lo, hi, condition):
    """condition(x)가 True가 되는 최소 x를 찾음 (lo~hi 범위)"""
    while lo < hi:
        mid = (lo + hi) // 2
        if condition(mid):
            hi = mid
        else:
            lo = mid + 1
    return lo

---

4. 정렬 커스텀 (key, 다중 기준)

가장 많이 검색되는 것 중 하나죠. sorted의 key에 람다를 넣으면 됩니다.

data = [("apple", 3), ("banana", 1), ("cherry", 2)]

# 두 번째 값 기준 오름차순
sorted(data, key=lambda x: x[1])

# 내림차순
sorted(data, key=lambda x: x[1], reverse=True)

# 다중 기준: 두 번째 값 오름차순, 같으면 첫 번째 값 내림차순
sorted(data, key=lambda x: (x[1], x[0]))

# 한 기준은 오름차순, 다른 기준은 내림차순일 때 (숫자라면 부호 반전)
sorted(data, key=lambda x: (x[1], -len(x[0])))

리스트 자체를 정렬하려면 .sort() (제자리 정렬, 반환값 없음):

arr.sort(key=lambda x: x[1])

---

5. collections (Counter, deque, defaultdict)

PS에서 체감 효율이 가장 큰 모듈입니다.

from collections import Counter, deque, defaultdict

# Counter: 개수 세기
c = Counter("aabbbc")          # {'b': 3, 'a': 2, 'c': 1}
c.most_common(2)               # 상위 2개 → [('b', 3), ('a', 2)]

# deque: 양방향 큐 (BFS의 핵심, O(1) popleft)
q = deque([1, 2, 3])
q.append(4)        # 오른쪽 추가
q.appendleft(0)    # 왼쪽 추가
q.pop()            # 오른쪽 제거
q.popleft()        # 왼쪽 제거 (list의 pop(0)보다 훨씬 빠름)

# defaultdict: 없는 키 접근 시 기본값 자동 생성
graph = defaultdict(list)
graph[1].append(2)             # 키 1이 없어도 에러 없이 [] 생성 후 추가

cnt = defaultdict(int)
cnt['x'] += 1                  # KeyError 없이 0에서 시작

---

6. 힙 / 우선순위 큐 (heapq)

파이썬 heapq는 최소 힙만 지원합니다. 최대 힙은 부호를 반전시키는 게 관용적인 방법입니다.

import heapq

# 최소 힙
heap = []
heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
heapq.heappush(heap, 2)
heapq.heappop(heap)            # → 1 (가장 작은 값)

# 최대 힙: 부호 반전
max_heap = []
heapq.heappush(max_heap, -3)
heapq.heappush(max_heap, -1)
-heapq.heappop(max_heap)       # → 3 (꺼낼 때 다시 반전)

# 리스트를 힙으로 한 번에 변환 (O(n))
arr = [5, 2, 8, 1]
heapq.heapify(arr)

# 튜플로 우선순위 지정: (우선순위, 데이터)
heapq.heappush(heap, (cost, node))

---

7. 메모이제이션 (lru_cache)

DP나 재귀 문제에서 캐싱을 직접 만들 필요 없이 데코레이터 한 줄로 끝냅니다.

from functools import lru_cache

@lru_cache(maxsize=None)       # 캐시 크기 무제한
def fib(n):
    if n < 2:
        return n
    return fib(n - 1) + fib(n - 2)

fib(100)                       # 캐싱 덕분에 즉시 계산

주의: lru_cache를 쓰는 함수의 인자는 해시 가능(immutable)해야 합니다. 리스트는 안 되고 튜플은 됩니다. Python 3.9+에서는 @cache로 더 간결하게 쓸 수 있습니다.

---

8. 순열 / 조합 / 곱집합 (itertools)

직접 백트래킹으로 짤 필요 없는 경우가 대부분입니다.

from itertools import permutations, combinations, product

arr = [1, 2, 3]

list(permutations(arr, 2))     # 순열: [(1,2),(1,3),(2,1),(2,3),(3,1),(3,2)]
list(combinations(arr, 2))     # 조합: [(1,2),(1,3),(2,3)]
list(product(arr, repeat=2))   # 중복순열: [(1,1),(1,2),...,(3,3)]

# 여러 리스트의 곱집합 (이중 for문 대체)
list(product([1, 2], ['a', 'b']))   # [(1,'a'),(1,'b'),(2,'a'),(2,'b')]

---

9. BFS 템플릿 (너비 우선 탐색)

격자/그래프 최단 거리의 표준 템플릿입니다.

from collections import deque

def bfs(start, graph, n):
    visited = [False] * (n + 1)
    q = deque([start])
    visited[start] = True

    while q:
        cur = q.popleft()
        for nxt in graph[cur]:
            if not visited[nxt]:
                visited[nxt] = True
                q.append(nxt)

2차원 격자 버전 (상하좌우 이동):

from collections import deque

def bfs_grid(grid, sr, sc):
    rows, cols = len(grid), len(grid[0])
    visited = [[False] * cols for _ in range(rows)]
    q = deque([(sr, sc)])
    visited[sr][sc] = True
    dist = [[0] * cols for _ in range(rows)]

    dr = [-1, 1, 0, 0]
    dc = [0, 0, -1, 1]

    while q:
        r, c = q.popleft()
        for d in range(4):
            nr, nc = r + dr[d], c + dc[d]
            if 0 <= nr < rows and 0 <= nc < cols and not visited[nr][nc]:
                visited[nr][nc] = True
                dist[nr][nc] = dist[r][c] + 1
                q.append((nr, nc))
    return dist

---

10. DFS 템플릿 (깊이 우선 탐색)

# 재귀 버전 (간결하지만 깊이 주의 → 2번 항목 참고)
def dfs(cur, graph, visited):
    visited[cur] = True
    for nxt in graph[cur]:
        if not visited[nxt]:
            dfs(nxt, graph, visited)

# 스택 버전 (재귀 한도 걱정 없음)
def dfs_iter(start, graph, n):
    visited = [False] * (n + 1)
    stack = [start]
    while stack:
        cur = stack.pop()
        if visited[cur]:
            continue
        visited[cur] = True
        for nxt in graph[cur]:
            if not visited[nxt]:
                stack.append(nxt)

---

11. 다익스트라 (최단 경로)

가중치가 있는 그래프의 최단 거리. 힙을 이용한 표준 구현입니다.

import heapq

def dijkstra(start, graph, n):
    """graph[u] = [(weight, v), ...], 노드 번호 1~n"""
    INF = float('inf')
    dist = [INF] * (n + 1)
    dist[start] = 0
    pq = [(0, start)]          # (누적 거리, 노드)

    while pq:
        cur_dist, cur = heapq.heappop(pq)
        if cur_dist > dist[cur]:   # 이미 더 짧은 경로로 처리됨
            continue
        for weight, nxt in graph[cur]:
            new_dist = cur_dist + weight
            if new_dist < dist[nxt]:
                dist[nxt] = new_dist
                heapq.heappush(pq, (new_dist, nxt))
    return dist

---

12. 유니온 파인드 (Union-Find / 분리 집합)

사이클 판별, 그래프 연결성, 크루스칼 알고리즘에 필수입니다.

def find(parent, x):
    if parent[x] != x:
        parent[x] = find(parent, parent[x])   # 경로 압축
    return parent[x]

def union(parent, a, b):
    ra, rb = find(parent, a), find(parent, b)
    if ra == rb:
        return False           # 이미 같은 집합 (사이클)
    parent[rb] = ra
    return True

# 초기화
n = 5
parent = [i for i in range(n + 1)]   # 0~n 자기 자신으로

---

13. 에라토스테네스의 체 (소수 구하기)

N 이하의 모든 소수를 빠르게 구합니다.

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

단일 숫자 소수 판별:

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

---

14. 최대공약수 / 최소공배수

직접 구현하지 말고 math 모듈을 쓰세요.

import math

math.gcd(12, 18)               # 최대공약수 → 6
math.lcm(4, 6)                 # 최소공배수 → 12 (Python 3.9+)

# 3.9 미만이라면 직접 계산
def lcm(a, b):
    return a * b // math.gcd(a, b)

---

15. 투 포인터 / 슬라이딩 윈도우

연속 구간 합, 부분 배열 문제의 단골 패턴입니다.

# 합이 target 이상이 되는 가장 짧은 구간 길이
def min_window(arr, target):
    left = 0
    total = 0
    res = float('inf')
    for right in range(len(arr)):
        total += arr[right]
        while total >= target:
            res = min(res, right - left + 1)
            total -= arr[left]
            left += 1
    return res if res != float('inf') else 0

---

16. 자주 하는 실수 모음

복붙만큼 자주 검색되는 게 "왜 안 되지?"입니다. 미리 피하세요.

2차원 배열 초기화 함정 — *로 복제하면 같은 리스트를 참조합니다.

# 잘못된 방법: 모든 행이 같은 리스트를 가리킴!
grid = [[0] * 3] * 3
grid[0][0] = 1                 # → 모든 행의 [0]이 1로 바뀜 (버그)

# 올바른 방법
grid = [[0] * 3 for _ in range(3)]

기본 정수 나눗셈 — /는 실수(float), 몫은 //입니다.

7 / 2      # → 3.5
7 // 2     # → 3
7 % 2      # → 1

무한대 값 — 거리 초기화에 자주 씁니다.

INF = float('inf')

---

마무리

여기 정리한 스니펫들은 대부분의 코딩테스트와 PS 문제의 70~80%를 커버합니다. 핵심은 "외우기"보다 "필요할 때 빠르게 꺼내 쓰기"입니다. 이 글을 북마크해두고, 문제를 풀다 막히는 부분이 생기면 해당 섹션만 복사해서 적용해보세요.

각 알고리즘의 원리를 더 깊이 이해하고 싶다면 따로 학습하는 게 좋지만, 실전에서는 검증된 템플릿을 손에 익히는 것만으로도 풀 수 있는 문제가 크게 늘어납니다.