Python 코딩테스트 치트시트

NANO·2일 전

[알고리즘 Cheat Sheet] 한 장 정리

목록 보기

10/10

코딩테스트에서 손에서 바로 나와야 하는 Python 문법/표준 라이브러리/패턴만 모았다. (Python 3.10 기준)

1. 기본 템플릿과 입출력

input()은 느리다... 입력이 1만 줄 넘으면 sys.stdin.readline으로 바꾸고, 재귀를 쓰면 setrecursionlimit을 올린다.

import sys
from collections import deque, defaultdict, Counter
from itertools import permutations, combinations, product, accumulate
from heapq import heappush, heappop, heapify
from bisect import bisect_left, bisect_right
from functools import lru_cache, cmp_to_key
import math

input = sys.stdin.readline
sys.setrecursionlimit(10**6)

n, m = map(int, input().split())
arr = list(map(int, input().split()))
grid = [list(map(int, input().split())) for _ in range(n)]
s = input().rstrip()          # readline은 개행 포함 → 문자열은 반드시 rstrip

# 출력 많으면 모아서 한 번에
out = []
for x in arr:
    out.append(str(x))
print("\n".join(out))
# 또는 sys.stdout.write("\n".join(out) + "\n")

상황	쓰는 것	비고
한 줄 정수 여러 개	`map(int, input().split())`	변수 개수 모르면 `list(...)`
문자열 그리드	`[input().rstrip() for _ in range(n)]`	각 행이 str, `grid[i][j]`로 접근
전체 입력 한 번에	`data = sys.stdin.read().split()` → 인덱스로 꺼내기	가장 빠름
테스트케이스 반복	`for _ in range(int(input())):`
숫자 ↔ 문자열	`int(s)`, `str(n)`, `float(s)`
진법 변환	`bin(n)[2:]`, `oct`, `hex`, `int("1011", 2)`, `format(n, 'b')`
포맷 출력	`f"{x:.2f}"`, `f"{n:05d}"`, `f"{s:>10}"`
구분자 출력	`print(arr)` (공백), `print(arr, sep="\n")`
끝 문자 지정	`print(x, end=" ")`

과제형(클래스 설계 등)이면 입력 파싱 대신 if __name__ == "__main__": 아래에 테스트 데이터를 만들어 직접 호출하고 print로 보여준다.

2. 문자열

Python 문자열도 불변이다. 반복 +=는 느리니 리스트에 모아 "".join().

할 일	코드
길이·인덱스·슬라이스	`len(s)`, `s[i]`, `s[a:b]` (b 미포함), `s[::-1]` (뒤집기), `s[::2]`
분리·합치기	`s.split()` (공백 전부), `s.split(",")`, `",".join(list)`
포함·위치·개수	`"ab" in s`, `s.find("ab")` (없으면 -1), `s.index("ab")` (없으면 에러), `s.count("a")`
바꾸기	`s.replace("a", "b")`, `s.replace("a", "b", 1)` (1회)
대소문자	`s.upper()`, `s.lower()`, `s.swapcase()`, `s.capitalize()`, `s.title()`
공백 제거	`s.strip()`, `s.lstrip()`, `s.rstrip()`
시작·끝	`s.startswith("ab")`, `s.endswith(("a", "b"))` (튜플 가능)
판별	`s.isdigit()`, `s.isalpha()`, `s.isalnum()`, `s.isupper()`, `s.islower()`, `s.isspace()`
문자 ↔ 코드	`ord('a')` → 97, `chr(97)` → 'a', `ord(c) - ord('a')` (알파벳 인덱스)
채우기	`s.zfill(5)`, `s.rjust(5, "0")`, `s.ljust(5)`, `s.center(5)`
정렬	`"".join(sorted(s))`, `sorted(s, reverse=True)`
반복	`"ab" * 3`
문자 교체 테이블	`s.translate(str.maketrans("abc", "xyz"))`
팰린드롬	`s == s[::-1]`
아스키 연산	`chr(ord(c) + 1)`

# 문자 빈도
from collections import Counter
cnt = Counter("banana")          # {'a': 3, 'n': 2, 'b': 1}
cnt.most_common(2)               # [('a', 3), ('n', 2)]

# 정규식은 간단한 것만
import re
re.sub(r"[^a-z0-9]", "", s)      # 영소문자·숫자만 남기기
re.findall(r"\d+", s)            # 숫자 덩어리 전부
re.split(r"[,;]", s)

3. 리스트·튜플·슬라이싱

리스트는 Java의 ArrayList, 튜플은 불변 리스트(해시 가능 → dict 키, set 원소로 쓸 수 있음).

a = [0] * n                        # 1차원 초기화
grid = [[0] * m for _ in range(n)] # 2차원 — [[0]*m]*n 은 같은 행 참조라 금지
visited = [[False] * m for _ in range(n)]

a.append(x);  a.insert(i, x);  a.pop();  a.pop(i);  a.remove(x)   # remove는 값, O(n)
a.index(x);  a.count(x);  a.reverse();  a.sort();  a.extend(b)
len(a);  sum(a);  max(a);  min(a);  sorted(a);  reversed(a)
a[-1]                              # 마지막
a[a.index(max(a))]                 # 최댓값 위치는 a.index(max(a))
b = a[:]  /  a.copy()  /  list(a)  # 얕은 복사
import copy; c = copy.deepcopy(grid)   # 2차원 복사

# 슬라이스
a[1:4];  a[:3];  a[3:];  a[::-1];  a[::2]
a[2:5] = [9, 9]                    # 구간 교체
del a[2:5]

# 컴프리헨션
sq = [x * x for x in a if x % 2 == 0]
flat = [x for row in grid for x in row]
transposed = list(zip(*grid))      # 전치 (튜플 행)
rotated = [list(r) for r in zip(*grid[::-1])]   # 시계방향 90도 회전

# 자주 쓰는 내장
for i, x in enumerate(a, start=1): ...
for x, y in zip(a, b): ...
any(x > 0 for x in a);  all(x > 0 for x in a)
list(map(str, a));  list(filter(None, a))
max(a, key=len);  max(d, key=d.get)   # 키 지정 최대

a.sort()는 제자리(None 반환), sorted(a)는 새 리스트. 리스트의 in은 O(n)이니 포함 확인이 많으면 set으로.

4. 딕셔너리·집합·collections

카운팅·그룹핑은 defaultdict와 Counter로 끝낸다. dict.get(k)는 없으면 None, d[k]는 KeyError.

# dict
d = {}
d[k] = v;  d.get(k, 0);  k in d;  del d[k];  d.pop(k, None)
d.keys();  d.values();  d.items()
for k, v in d.items(): ...
sorted(d.items(), key=lambda kv: -kv[1])       # 값 내림차순
dict(sorted(d.items()))                        # 키 정렬된 새 dict
d.setdefault(k, []).append(v)                  # 그룹핑 (defaultdict 없이)
# dict는 3.7+ 삽입 순서 유지 → LinkedHashMap 역할

# defaultdict / Counter
from collections import defaultdict, Counter
cnt = defaultdict(int);   cnt[k] += 1
grp = defaultdict(list);  grp[k].append(v)
c = Counter(arr)                               # 빈도
c.most_common(3);  c[x];  c.total()  (3.10+)
Counter(a) - Counter(b);  Counter(a) & Counter(b)   # 차집합·교집합 (빈도 기준)

# set
s = set();  s.add(x);  s.discard(x) (없어도 OK);  s.remove(x) (없으면 에러);  x in s
a | b;  a & b;  a - b;  a ^ b                  # 합·교·차·대칭차
len(set(arr))                                  # 중복 제거 개수
frozenset(...)                                 # 해시 가능한 set (set의 set 만들 때)

# deque (큐·스택·슬라이딩)
from collections import deque
q = deque([1, 2])
q.append(x);  q.appendleft(x);  q.pop();  q.popleft();  q[0];  q[-1]
q.rotate(1)                                    # 오른쪽으로 한 칸
deque(maxlen=k)                                # 고정 길이 윈도우
# 리스트 pop(0)은 O(n) — 큐는 반드시 deque

# heapq (최소 힙)
import heapq
h = [];  heapq.heappush(h, x);  heapq.heappop(h);  h[0]
heapq.heapify(arr)                             # 리스트를 제자리 힙으로 O(n)
heapq.heappush(h, -x)                          # 최대 힙은 부호 뒤집기
heapq.heappush(h, (priority, data))            # 튜플은 첫 원소 기준
heapq.nlargest(3, arr);  heapq.nsmallest(3, arr, key=...)

# 정렬 유지 삽입 / 이분탐색
from bisect import bisect_left, bisect_right, insort
bisect_left(a, x)    # x 이상인 첫 인덱스 (lower bound)
bisect_right(a, x)   # x 초과인 첫 인덱스 (upper bound) → 개수 = right - left
insort(a, x)         # 정렬 유지 삽입 O(n)

자료구조	언제	비용
list	순서 목록, 인덱스 접근, 스택	끝 append/pop O(1), 앞 삽입/삭제 O(n), in O(n)
dict / set	카운팅, 중복 체크, 조회	O(1)
deque	큐, BFS, 양끝 조작	O(1)
heapq	최소/최대 반복 추출, 다익스트라	O(log n)
Counter / defaultdict	빈도, 그룹핑	O(1)
bisect + 정렬 리스트	범위 검색 (TreeMap 대용)	검색 O(log n), 삽입 O(n)

Python엔 TreeMap이 없다. 정렬 상태 유지 + 범위 검색이 필요하면 정렬 리스트 + bisect, 또는 힙으로 우회.

5. 정렬과 key

sort/sorted는 안정 정렬(TimSort). 다중 조건은 key에 튜플, 내림차순은 reverse=True 또는 숫자면 부호 반전.

a.sort()                                  # 오름차순 제자리
sorted(a, reverse=True)                   # 내림차순 새 리스트
sorted(words, key=len)                    # 길이 순
sorted(words, key=lambda w: (len(w), w))  # 길이 순, 같으면 사전순
sorted(people, key=lambda p: (-p[1], p[0]))   # 점수 내림, 이름 오름
sorted(d.items(), key=lambda kv: (-kv[1], kv[0]))

# 문자열 내림 + 숫자 오름처럼 부호 반전이 안 되는 경우 → cmp_to_key
from functools import cmp_to_key
def cmp(x, y):
    if x[0] != y[0]: return -1 if x[0] > y[0] else 1      # 문자열 내림
    return x[1] - y[1]                                    # 숫자 오름
sorted(a, key=cmp_to_key(cmp))

# 가장 큰 수 만들기 (문자열 결합 비교)
sorted(nums, key=cmp_to_key(lambda a, b: 1 if a + b < b + a else -1))

# 객체 정렬
sorted(songs, key=lambda s: (-s.plays, s.title))
from operator import itemgetter, attrgetter
sorted(rows, key=itemgetter(1, 0));  sorted(songs, key=attrgetter("plays"))

key 함수는 원소당 한 번만 호출돼서 cmp_to_key보다 빠르다. 가능하면 튜플 key로 해결하고 cmp는 정말 필요할 때만.

6. itertools·functools·math

조합·순열은 직접 짜지 말고 itertools. 메모이제이션은 lru_cache 한 줄.

from itertools import permutations, combinations, product, combinations_with_replacement, accumulate, groupby, chain

list(permutations([1,2,3], 2))     # 순열 (순서 있음) nP2
list(combinations([1,2,3], 2))     # 조합 (순서 없음) nC2
list(product([0,1], repeat=3))     # 중복 순열 / 데카르트 곱 → 비트마스크 대용
list(combinations_with_replacement([1,2,3], 2))   # 중복 조합
list(accumulate([1,2,3]))          # 누적합 [1,3,6]
list(accumulate(a, max))           # 누적 최댓값
list(chain(a, b))                  # 이어붙이기
for k, g in groupby(sorted(a)): print(k, len(list(g)))   # 연속 동일값 묶기 (정렬 먼저)

from functools import lru_cache, reduce
@lru_cache(maxsize=None)
def fib(n):
    return n if n < 2 else fib(n-1) + fib(n-2)
# 3.9+ 는 @cache 도 가능. 인자는 해시 가능해야 함 (리스트 → 튜플)
reduce(lambda x, y: x * y, a, 1)

import math
math.gcd(a, b);  math.lcm(a, b) (3.9+);  math.gcd(*arr) (3.9+)
math.factorial(n);  math.comb(n, r);  math.perm(n, r)
math.ceil(x);  math.floor(x);  math.sqrt(x);  math.isqrt(n) (정수 제곱근)
math.inf;  -math.inf;  float("inf")
math.log(x, base);  math.log2(x);  math.hypot(dx, dy);  math.dist(p, q)
divmod(a, b)       # (몫, 나머지)
pow(a, b, mod)     # 모듈러 거듭제곱 O(log b)
abs(x);  round(x, 2)   # round는 은행가 반올림 — 0.5는 짝수로

7. 자주 나오는 알고리즘 패턴

완전탐색

# 비트마스크 부분집합
for mask in range(1 << n):
    subset = [a[i] for i in range(n) if mask & (1 << i)]

# 백트래킹 (N-Queen 틀)
def dfs(depth, path):
    if depth == n:
        result.append(path[:])
        return
    for i in range(n):
        if not used[i]:
            used[i] = True;  path.append(i)
            dfs(depth + 1, path)
            path.pop();  used[i] = False

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

left = total = 0;  best = float("inf")
for right in range(n):
    total += a[right]
    while total >= target:
        best = min(best, right - left + 1)
        total -= a[left];  left += 1

누적합

pre = [0] + list(accumulate(a))     # pre[i] = a[0..i-1] 합
# [l, r] 구간 합 = pre[r+1] - pre[l]

# 2차원
P = [[0] * (m + 1) for _ in range(n + 1)]
for i in range(n):
    for j in range(m):
        P[i+1][j+1] = grid[i][j] + P[i][j+1] + P[i+1][j] - P[i][j]

BFS (그리드 최단거리)

from collections import deque
dr, dc = [-1, 1, 0, 0], [0, 0, -1, 1]
dist = [[-1] * m for _ in range(n)]
q = deque([(sr, sc)]);  dist[sr][sc] = 0
while q:
    r, c = q.popleft()
    for d in range(4):
        nr, nc = r + dr[d], c + dc[d]
        if not (0 <= nr < n and 0 <= nc < m): continue
        if grid[nr][nc] == 1 or dist[nr][nc] != -1: continue
        dist[nr][nc] = dist[r][c] + 1
        q.append((nr, nc))

DFS (인접 리스트, 재귀 → 스택)

adj = [[] for _ in range(n)]
adj[a].append(b);  adj[b].append(a)

visited = [False] * n
def dfs(u):
    visited[u] = True
    for v in adj[u]:
        if not visited[v]: dfs(v)

# 재귀 깊이 걱정되면 스택으로
stack = [start];  visited[start] = True
while stack:
    u = stack.pop()
    for v in adj[u]:
        if not visited[v]:
            visited[v] = True;  stack.append(v)

다익스트라

import heapq
INF = float("inf")
dist = [INF] * n;  dist[s] = 0
pq = [(0, s)]
while pq:
    d, u = heapq.heappop(pq)
    if d > dist[u]: continue
    for v, w in adj[u]:
        if dist[u] + w < dist[v]:
            dist[v] = dist[u] + w
            heapq.heappush(pq, (dist[v], v))

이분탐색 / 매개변수 탐색

from bisect import bisect_left
idx = bisect_left(a, target)          # 정렬된 a에서 target 이상 첫 위치

# "X가 가능한가"가 단조일 때 최대 X
lo, hi = 1, max_x
while lo < hi:
    mid = (lo + hi + 1) // 2
    if ok(mid): lo = mid
    else: hi = mid - 1
# 최소 X 는 mid = (lo+hi)//2, ok면 hi = mid, 아니면 lo = mid+1

유니온 파인드

parent = list(range(n))
def find(x):
    while parent[x] != x:
        parent[x] = parent[parent[x]];  x = parent[x]
    return x
def union(a, b):
    parent[find(a)] = find(b)

DP 기본 틀

# 1차원
dp = [0] * (n + 1)
for i in range(1, n + 1):
    dp[i] = max(dp[i-1], ...)

# 2차원 (LCS, 격자 경로)
dp = [[0] * (m + 1) for _ in range(n + 1)]

# 탑다운
@lru_cache(None)
def solve(i, j): ...

# 배낭
for w, v in items:
    for cap in range(W, w - 1, -1):      # 0/1은 역순
        dp[cap] = max(dp[cap], dp[cap - w] + v)

8. 자주 하는 실수와 함정

함정	틀린 예	맞는 예
2차원 초기화	`[[0]m]n` → 모든 행이 같은 객체	`[[0]*m for _ in range(n)]`
readline 개행	`s = input()` 후 `len(s)`에 개행 포함	`input().rstrip()`
리스트 복사	`b = a` (같은 참조)	`b = a[:]`, 2차원은 `deepcopy`
순회 중 삭제	`for x in a: a.remove(x)`	`a = [x for x in a if cond]`
큐를 list로	`a.pop(0)` O(n)	`deque.popleft()`
`sort` 반환값	`a = a.sort()` → None	`a.sort()` 또는 `a = sorted(a)`
정수 나눗셈	`7 / 2` → 3.5 (float)	`7 // 2` → 3
음수 나눗셈	`-7 // 2` → -4 (바닥)	0 방향 자르기는 `int(-7 / 2)` → -3
큰 수 나눗셈 float	`10**18 / 3` 정밀도 손실	`//` 유지, float 변환 피하기
재귀 깊이	기본 1000 → RecursionError	`sys.setrecursionlimit(10**6)` + 깊이 10만 넘으면 스택으로
전역 변수 수정	함수 안 `cnt += 1` → UnboundLocalError	`global cnt` 또는 리스트로 감싸기 `cnt[0] += 1`
가변 기본 인자	`def f(a, acc=[])`	`def f(a, acc=None): acc = acc or []`
dict 키 없음	`d[k]` KeyError	`d.get(k, 0)` 또는 `defaultdict`
set 원소로 리스트	`{[1,2]}` TypeError	`{(1,2)}` 튜플
lru_cache 인자	리스트 넘김	튜플/문자열로 변환
`is` vs `==`	`x is 1000`	`x == 1000` (`is`는 None 비교에만)
문자열 결합 반복	`s += c` 루프	`parts.append(c)` 후 `"".join(parts)`
`max([])`	빈 리스트 → ValueError	`max(a, default=0)`
인덱스 음수	`a[i-1]`에서 i=0 → 마지막 원소	경계 검사 먼저
부동소수 비교	`0.1 + 0.2 == 0.3` False	`abs(x - y) < 1e-9` 또는 정수로 스케일
round	`round(2.5)` → 2	올림 반올림은 `int(x + 0.5)`
입력 속도	`input()` 10만 번	`sys.stdin.readline`
출력 속도	`print` 10만 번	리스트에 모아 `"\n".join`
시간 초과 일반	1초 ≈ 2천만 연산 기준으로 복잡도 계산 안 함	n=10^5면 O(n log n)까지, n=10^3이면 O(n²) OK

9. 속도 올리는 습관

Python은 느리다. 같은 O(n)이라도 쓰는 방법에 따라 3~10배 차이 난다.

for 안에서 len(a), a.b.c 같은 속성 접근 반복하지 말고 변수로 빼기
in 검사는 list 말고 set/dict
sum, max, min, sorted, "".join, 컴프리헨션은 C로 돌아서 명시적 루프보다 빠름
2차원 접근 grid[i][j] 반복이면 row = grid[i] 한 번 꺼내기
재귀 DP보다 반복 DP가 2~3배 빠름. 깊이 깊으면 반복으로
전역 변수보다 지역 변수가 빠름 → 메인 로직을 def main(): 안에 넣고 호출
PyPy 선택 가능하면 PyPy (단, 재귀와 큰 dict에서 메모리 주의)
문자열 비교 많으면 ord로 정수화

def main():
    input = sys.stdin.readline
    # ... 전체 로직 ...
    print(...)

if __name__ == "__main__":
    main()

10. 과제형 테스트용 클래스 설계

오프라인 IDE 테스트에서 "플레이리스트 관리 기능을 만들어라" 같은 과제가 나오면 이 틀로 간다.

from dataclasses import dataclass, field
from typing import Optional

@dataclass
class Song:
    title: str
    artist: str
    duration: int          # 초
    plays: int = 0

class Playlist:
    def __init__(self, name: str):
        self.name = name
        self._songs: list[Song] = []
        self._titles: set[str] = set()     # 중복 방지 O(1)

    def add(self, song: Song) -> bool:
        if song.title in self._titles:
            return False
        self._songs.append(song)
        self._titles.add(song.title)
        return True

    def remove(self, title: str) -> Optional[Song]:
        for i, s in enumerate(self._songs):
            if s.title == title:
                self._titles.discard(title)
                return self._songs.pop(i)
        return None

    def by_artist(self, artist: str) -> list[Song]:
        return [s for s in self._songs if s.artist == artist]

    def total_duration(self) -> int:
        return sum(s.duration for s in self._songs)

    def top(self, n: int) -> list[Song]:
        return sorted(self._songs, key=lambda s: -s.plays)[:n]

    def __len__(self): return len(self._songs)
    def __repr__(self): return f"Playlist({self.name}, {len(self)} songs)"


if __name__ == "__main__":
    pl = Playlist("드라이브")
    pl.add(Song("A", "IU", 210, 50))
    pl.add(Song("B", "IU", 180, 120))
    print(pl.add(Song("A", "IU", 210)))   # False — 중복
    print(pl.by_artist("IU"))
    print(pl.total_duration())
    print(pl.top(1))

면접에서 받을 질문과 한 줄 답:
왜 set을 따로 뒀나 → 중복 체크를 O(1)로 하려고, 리스트 in은 O(n)
곡이 100만 개면 → 삭제가 O(n)이니 dict로 title→Song 매핑하고 순서가 필요하면 OrderedDict나 인덱스 관리
dataclass 쓴 이유 → __init__, __repr__, __eq__ 자동 생성으로 보일러플레이트 제거
타입 힌트는 → 실행엔 영향 없지만 의도를 읽는 사람에게 보여준다.
변수명은 a, tmp 대신 play_count, songs_by_artist처럼
Python은 snake_case

NANO

즐거운 토마토

이전 포스트