PS를 위한 파이썬 정리
코딩테스트에서 자주 사용했던 파이썬 문법, 라이브러리, 알고리즘 템플릿 등 정리
1. 입출력
기본 입력
input()은 한 줄을 문자열로 읽고, split()은 공백 기준으로 나눈다. map()으로 형변환하면 정수 리스트를 만들 수 있다.
# 한 줄 읽기
s = input()
# 공백으로 구분된 입력을 리스트로
arr = list(map(int, input().split()))
# 여러 개의 값
n, m = map(int, input().split())
# 정수 리스트
line = list(map(int, input().split()))
2차원 리스트 입력
행마다 입력을 받아 2차원 리스트를 구성한다. 리스트 컴프리헨션으로 한 줄로 줄일 수 있다.
# n개 행, m개 열
n, m = map(int, input().split())
grid = []
for _ in range(n):
row = list(map(int, input().split()))
grid.append(row)
# 리스트 컴프리헨션 버전, 위의 반복문과 동일하다.
grid = [list(map(int, input().split())) for _ in range(n)]
빠른 입력 (sys.stdin)
시간 제한이 타이트할 때는 sys.stdin.readline() 을 사용한다.
import sys
input = sys.stdin.readline
n = int(input())
arr = list(map(int, input().split()))
리스트 출력
*arr로 언팩킹하면 공백 구분으로 한 줄 출력된다. 세로 출력이 필요하면 join을 쓴다.
arr = [1, 2, 3, 4, 5]
# 방법 1: 언팩킹
print(*arr) # 1 2 3 4 5
# 방법 2: join (문자열)
print('\n'.join(map(str, arr))) # 1\n2\n3\n4\n5
# 방법 3: 루프
for x in arr:
print(x)
빠른 출력
join은 구분자를 지정해서 문자열을 합칠 때 쓴다. 출력이 많을 때는 모아서 한 번에 출력하면 시간초과를 피할 수 있다.
import sys
result = []
for _ in range(n):
result.append(str(answer))
print('\n'.join(result))
2. 자료형 비교
| 자료형 | 문법 | 가변 | 순서 | 중복 | 조회 | 수정 | 비고 |
|---|---|---|---|---|---|---|---|
| List | [1, 2, 3] | O | O | O | O(1) | O | 동적 배열 |
| Tuple | (1, 2, 3) | X | O | O | O(1) | X | 해시 가능 |
| Dict | {k: v} | O | O (3.7+) | X (key) | O(1) | O | key-value |
| Set | {1, 2, 3} | O | X | X | O(1) | O | 중복 제거 |
시간복잡도 정리
| 연산 | List | Dict | Set | 설명 |
|---|---|---|---|---|
| 조회 | O(n) | O(1) | O(1) | Dict/Set은 해시 사용 |
| 삽입 | O(n) | O(1) | O(1) | List 맨 앞 삽입은 O(n) |
| 삭제 | O(n) | O(1) | O(1) | 위치를 알아야 함 |
| 정렬 | O(n log n) | - | - | 비교 기반 |
in 연산 | O(n) | O(1) | O(1) | Set/Dict이 훨씬 빠름 |
주요 메서드
자주 사용하는 메서드를 자료형별로 정리했다. 자료형에 따라 시간복잡도가 다르므로 주의할 필요가 있다.
List
arr.append(x) # 맨 뒤에 추가
arr.extend([4, 5]) # 여러 원소 추가
arr.pop() # 맨 뒤 제거, 반환
arr.pop(0) # 첫 원소 제거 (O(n))
arr.insert(i, x) # i번 위치에 삽입 (O(n))
arr.remove(x) # 첫 x 제거 (O(n))
arr.reverse() # 역순
arr.sort() # 정렬 (인자 가능)
arr.count(x) # x의 개수
arr.index(x) # x의 위치
Dict
d[key] = value # 할당
d.get(key, default) # key 없으면 default 반환
d.pop(key) # key 제거, 값 반환
d.keys() # 모든 key
d.values() # 모든 value
d.items() # (key, value) 쌍
d.update(other_dict) # 다른 딕셔너리로 업데이트
Set
s.add(x) # 추가
s.remove(x) # 제거 (없으면 에러)
s.discard(x) # 제거 (없으면 무시)
s.pop() # 임의의 원소 제거
s & t # 교집합
s | t # 합집합
s - t # 차집합
s ^ t # 대칭차 (둘 중 하나만)
Truthy와 Falsy
빈 컬렉션, 0, None은 False로 평가된다. 스택/큐가 비었는지 확인할 때 if stk: 형태로 사용한다.
# Falsy: 0, None, [], {}, "", (), False
if not x: # x가 거짓 값이면
pass
# Truthy: 0이 아닌 수, 비어있지 않은 컨테이너, True
if x: # x가 참 값이면
pass
삼항 연산자
if-else를 한 줄로 쓰는 문법이다. 스택 문제에서 "비었으면 -1, 아니면 pop" 패턴에 자주 쓴다.
x = 10 if condition else 20
for-else
for 루프가 break 없이 끝까지 돌면 else가 실행된다. 검색 실패 처리에 유용하다. (예: 그룹 단어 체커)
arr = [1, 2, 3, 4, 6, 7]
for x in arr:
if x == 5:
print("found")
break
else:
print("not found")
- 5를 찾으면
break되어 else 실행이 안 된다. - 끝까지 못 찾으면
break되지 않아 else가 실행된다.
3. 리스트 컴프리헨션
기본 문법
반복문으로 리스트를 생성하는 축약 문법이다. if만 쓸 때는 for 뒤에, if-else는 for 앞에 온다.
# 1차원
squares = [x**2 for x in range(10)]
# 조건부
evens = [x for x in range(10) if x % 2 == 0]
# 중첩
pairs = [(x, y) for x in range(3) for y in range(3)]
# 조건부 (if-else)
result = [x if x % 2 == 0 else -x for x in range(5)] # [0, -1, 2, -3, 4]
고정 크기 배열 초기화
[0]*n은 1차원에서는 안전하지만, [[0]*m]*n은 모든 행이 같은 객체를 참조하므로 컴프리헨션을 써야 한다.
arr = [0] * 5 # 1D
grid = [[0] * 4 for _ in range(3)] # 2D (올바른 방법)
grid = [[0] * 4] * 3 # 2D (X: 같은 행 참조)
2차원 리스트 깊은 복사
2차원 리스트를 =로 복사하면 같은 객체를 참조한다. [row[:] for row in grid]로 독립적인 복사본을 만든다. 브루트포스에서 매번 원본을 유지한 채 시뮬레이션해야 할 때 자주 쓴다. (예: 연구소 문제처럼 매 조합마다 새 지도가 필요한 경우)
import copy
# 방법 1: copy 라이브러리
grid2 = copy.deepcopy(grid)
# 방법 2: 리스트 컴프리헨션
grid2 = [row[:] for row in grid]
2차원 격자에서 좌표
이중 for문으로 격자의 각 칸을 행·열 인덱스와 함께 순회한다.
for i in range(len(grid)):
for j in range(len(grid[i])):
value = grid[i][j]
4. lambda와 정렬
lambda 기본
이름 없는 익명 함수를 만든다. 정렬의 key 파라미터에 주로 사용한다.
square = lambda x: x**2
print(square(5)) # 25
add = lambda x, y: x + y
print(add(3, 4)) # 7
sorted() vs sort()
sorted()는 새 리스트를 반환하고, sort()는 원본을 수정하며 None을 반환한다.
arr = [3, 1, 4, 1, 5]
# sorted(): 새로운 정렬된 리스트 반환
result = sorted(arr) # [1, 1, 3, 4, 5]
# sort(): 리스트 내부에서 정렬 (반환값 없음)
arr.sort() # arr = [1, 1, 3, 4, 5]
key 파라미터
정렬 기준을 지정한다. lambda와 함께 특정 원소나 조건을 기준으로 정렬할 수 있다.
# 절댓값
arr = [3, -1, -4, 2]
sorted(arr, key=abs) # [-1, 2, 3, -4]
# 문자열 길이
words = ["apple", "pie", "a"]
sorted(words, key=len) # ["a", "pie", "apple"]
# 두 번째 원소
pairs = [(1, 3), (1, 1), (2, 2)]
sorted(pairs, key=lambda x: x[1]) # [(1, 1), (2, 2), (1, 3)]
다중 조건 정렬
튜플을 반환하면 앞 원소부터 우선 정렬된다. 내림차순은 숫자에 -를 붙인다. (예: 나이순 정렬)
# 첫 번째 요소는 오름차순, 두 번째는 내림차순
students = [("Alice", 85), ("Bob", 85), ("Charlie", 90)]
sorted(students, key=lambda x: (x[1], -ord(x[0][0])))
# 또는 방법 2: 여러 번 sort() (역순으로, stable sort)
students.sort(key=lambda x: x[0]) # 이름순
students.sort(key=lambda x: x[1], reverse=True) # 점수 내림차순
sort() 체이닝은 불가능하다. (sort는 None 반환)
arr.sort().sort() # X - 에러
5. 내장함수
sum, max, min
기본 집계 함수다. key 파라미터로 기준을 바꿀 수 있다.
arr = [1, 3, 2, 5, 4]
sum(arr) # 15
max(arr) # 5
min(arr) # 1
words = ["a", "abc", "ab"]
max(words, key=len) # "abc"
sum(arr, 10) # 25 (10부터 시작)
enumerate
인덱스와 값을 동시에 순회한다. start 파라미터로 시작 인덱스를 지정할 수 있다.
for i, val in enumerate(arr):
print(i, val) # (0, 1), (1, 3), (2, 2), ...
for i, val in enumerate(arr, start=1):
print(i, val) # (1, 1), (2, 3), ...
reversed
리스트를 뒤에서부터 순회한다. [::-1]과 달리 새 리스트를 만들지 않아 메모리 효율적이다.
for val in reversed(arr):
print(val)
# 리스트로 변환
rev_arr = list(reversed(arr))
any / all
any()는 하나라도 참이면 True, all()은 모두 참이어야 True다.
arr = [1, 2, 3]
any(arr) # True
all(arr) # True
arr = [0, 1, 2]
any(arr) # True
all(arr) # False
any(x > 5 for x in arr) # 5보다 큰 값이 있나?
all(x > 0 for x in arr) # 모두 양수인가?
divmod
몫과 나머지를 동시에 반환한다. 시간 변환 문제에서 유용하다. (예: 시각)
quotient, remainder = divmod(17, 5) # (3, 2)
zip
여러 iterable을 병렬로 순회한다. 행렬 전치에도 활용한다.
a = [1, 2, 3]
b = [4, 5, 6]
for x, y in zip(a, b):
print(x, y) # (1, 4), (2, 5), (3, 6)
# 행렬 전치
matrix = [[1, 2], [3, 4], [5, 6]]
transposed = list(zip(*matrix)) # [(1, 3, 5), (2, 4, 6)]
map / filter
map()은 각 원소에 함수를 적용하고, filter()는 조건에 맞는 원소만 추출한다.
nums = list(map(int, ["1", "2", "3"])) # [1, 2, 3]
evens = list(filter(lambda x: x % 2 == 0, range(10))) # [0, 2, 4, 6, 8]
6. 자주 쓰는 라이브러리
collections.deque
양쪽 끝에서 O(1) 시간에 추가/제거가 가능한 큐(Queue) 구조다. BFS에서 자주 사용한다. (예: 덱)
from collections import deque
q = deque()
q.append(1) # 오른쪽에 추가
q.appendleft(0) # 왼쪽에 추가
q.pop() # 오른쪽 제거, 반환
q.popleft() # 왼쪽 제거, 반환
q[0] # 첫 원소 접근
collections.Counter
원소의 개수를 세서 딕셔너리로 반환한다. (예: 숫자 카드 2)
from collections import Counter
s = "aabbcc"
c = Counter(s) # Counter({'a': 2, 'b': 2, 'c': 2})
c['a'] # 2
c.most_common(2) # [('a', 2), ('b', 2)]
collections.defaultdict
존재하지 않는 key에 접근하면 기본값을 자동 생성한다.
from collections import defaultdict
d = defaultdict(int)
d['a'] += 1 # key 없어도 0으로 초기화 후 1 추가
d = defaultdict(list)
d['group'].append(1) # key 없어도 빈 리스트로 초기화
heapq (최소 힙)
파이썬은 기본적으로 최소 힙을 제공한다. 최대 힙은 값에 -를 붙여서 구현한다. 이미 리스트가 있을 때는 heapify()로 O(n)에 힙으로 변환할 수 있다.
import heapq
heap = []
heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
heapq.heappush(heap, 2)
heapq.heappop(heap) # 1
heapq.heappop(heap) # 2
nums = [3, 1, 2]
heapq.heapify(nums) # O(n)
최대 힙 구현 (음수 값 저장)
heap = []
heapq.heappush(heap, -3)
heapq.heappush(heap, -1)
val = -heapq.heappop(heap) # 3
itertools (순열, 조합, 중복순열, 중복조합)
순서를 고려하면 순열(permutations), 순서 상관없이 고르면 조합(combinations), 같은 원소를 여러 번 뽑을 수 있으면 중복순열(product) 또는 중복조합(combinations_with_replacement)을 사용한다.
| 유형 | 순서 | 중복 | 함수 |
|---|---|---|---|
| 순열 | O | X | permutations |
| 조합 | X | X | combinations |
| 중복순열 | O | O | product |
| 중복조합 | X | O | combinations_with_replacement |
순열 (순서 O, 중복 X) — N과 M (1)
from itertools import permutations
arr = [1, 2, 3]
for perm in permutations(arr):
print(perm) # (1,2,3), (1,3,2), (2,1,3), ...
# r개 선택 (순열)
for perm in permutations(arr, 2):
print(perm) # (1,2), (1,3), (2,1), (2,3), (3,1), (3,2)
조합 (순서 X, 중복 X) — N과 M (2)
from itertools import combinations
arr = [1, 2, 3]
for comb in combinations(arr, 2):
print(comb) # (1,2), (1,3), (2,3)
중복순열 (순서 O, 중복 O) — N과 M (3)
from itertools import product
for p in product([1, 2], ['a', 'b']):
print(p) # (1,'a'), (1,'b'), (2,'a'), (2,'b')
# repeat 파라미터로 자신과의 곱
for p in product([1, 2], repeat=2):
print(p) # (1,1), (1,2), (2,1), (2,2)
중복조합 (순서 X, 중복 O) — N과 M (4)
from itertools import combinations_with_replacement
arr = [1, 2, 3]
for comb in combinations_with_replacement(arr, 2):
print(comb) # (1,1), (1,2), (1,3), (2,2), (2,3), (3,3)
단순히 모든 경우를 나열하기만 하면 itertools가 간편하다. 하지만 "합이 10 이하인 조합만" 같은 조건이 있으면 itertools는 전부 생성한 뒤 필터링해야 하므로 느리다. 이럴 때는 백트래킹으로 직접 구현하면 조건에 안 맞는 경우를 중간에 잘라낼 수 있어서 훨씬 빠르다.
bisect (이진 탐색)
정렬된 리스트에서 이진 탐색으로 삽입 위치를 찾는다.
import bisect
arr = [1, 3, 5, 7, 9]
# bisect_left: x가 들어갈 위치 (x 미만인 요소 개수)
bisect.bisect_left(arr, 5) # 2
bisect.bisect_left(arr, 4) # 2
# bisect_right: x 이하인 요소 개수
bisect.bisect_right(arr, 5) # 3
bisect.bisect_right(arr, 4) # 2
# insort: 정렬 상태 유지하며 삽입
bisect.insort(arr, 6) # arr = [1, 3, 5, 6, 7, 9]
정렬된 배열에서 특정 범위에 값이 몇 개 있는지 O(log n)에 구할 수 있다. for문으로 세면 O(n)이다.
# arr에서 4 이상 7 이하인 값의 개수
left = bisect.bisect_left(arr, 4)
right = bisect.bisect_right(arr, 7)
count = right - left
math
import math
math.gcd(12, 8) # 4
math.lcm(12, 8) # 24
math.comb(5, 2) # 10 (5C2)
math.perm(5, 2) # 20 (5P2)
math.factorial(5) # 120
math.sqrt(16) # 4.0
math.ceil(3.2) # 4
math.floor(3.8) # 3
math.isqrt(17) # 4 (정수 제곱근)
functools.lru_cache
재귀로 DP를 짤 때(탑다운), 같은 인자로 다시 호출되면 저장된 결과를 반환해주는 데코레이터다. 직접 memo 딕셔너리를 관리하지 않아도 된다.
lru_cache 없이 직접 메모이제이션하면 매번 딕셔너리를 관리해야 한다.
memo = {}
def fib(n):
if n in memo: # 이미 계산한 적 있으면
return memo[n] # 저장된 값 반환
if n < 2:
return n
memo[n] = fib(n-1) + fib(n-2) # 계산하고 저장
return memo[n]
@lru_cache 한 줄을 붙이면 위와 동일하게 동작한다. 자동으로 결과를 캐싱해준다.
from functools import lru_cache
@lru_cache(maxsize=None)
def fib(n):
if n < 2:
return n
return fib(n-1) + fib(n-2)
7. 문자열/숫자 처리 패턴
정수 나눗셈
정수끼리 나눌 때는 //를 사용한다. /는 결과가 float이 되어 큰 수에서 정밀도가 깨질 수 있다.
17 // 5 # 3 (몫)
17 / 5 # 3.4 (실수)
17 % 5 # 2 (나머지)
swap
파이썬에서는 별도의 임시 변수 없이 한 줄로 값을 교환할 수 있다.
a, b = b, a # 파이썬 스타일
ASCII 값
ord()는 문자를 아스키코드 숫자로, chr()는 숫자를 문자로 변환한다. (예: 알파벳 찾기)
ord('A') # 65
chr(65) # 'A'
ord('a') - ord('A') # 32 (대소문자 차이)
슬라이싱과 불변성
[start:end:step] 형태로 부분 리스트/문자열을 추출한다. 문자열은 불변이므로 수정하려면 리스트로 변환해야 한다.
s = "hello"
s[1:3] # "el"
s[::-1] # "olleh" (역순)
s.replace('l', 'L') # "heLLo" (새 문자열)
# 문자열은 불변이므로 s[0] = 'H' 불가능
# 리스트로 변환 후 다시 문자열로
chars = list(s)
chars[0] = 'H'
s = ''.join(chars)
문자열 <-> 리스트
s = "123"
arr = list(s) # ['1', '2', '3']
s = ''.join(arr) # "123"
# 문자가 아닌 경우
nums = [1, 2, 3]
s = ''.join(map(str, nums)) # "123"
빈도 세기 (3가지 방법)
같은 결과를 3가지 방법으로 구현할 수 있다. (예: 숫자 카드 2)
from collections import Counter, defaultdict
s = "aabbcc"
# 방법 1: Counter
freq = Counter(s)
freq['a'] # 2
# 방법 2: dict.get()
freq = {}
for c in s:
freq[c] = freq.get(c, 0) + 1
# 방법 3: defaultdict
freq = defaultdict(int)
for c in s:
freq[c] += 1
좌표를 Set에 저장
BFS/DFS에서 2차원 좌표의 방문 여부를 체크할 때, 튜플로 Set에 저장하면 O(1)로 확인 가능하다.
visited = set()
visited.add((0, 0))
if (0, 0) in visited:
print("이미 방문")
무한대
최솟값/최댓값을 찾을 때 초기값으로 사용한다.
float('inf') # 양의 무한대
-float('inf') # 음의 무한대
min_val = float('inf')
min_val = min(min_val, 10) # 10
중복 제거 후 정렬
set으로 중복을 제거하고 sorted로 정렬하는 한 줄 패턴이다. (예: 좌표 압축)
arr = [3, 1, 2, 1, 3]
result = sorted(set(arr)) # [1, 2, 3]
두 구간의 겹치는 범위
left = max(a1, a2)
right = min(b1, b2)
if left <= right:
# [left, right]가 겹치는 부분
length = right - left + 1
8. 투 포인터
두 개의 포인터(인덱스)를 움직이며 조건을 만족하는 구간이나 쌍을 찾는 기법이다.
구간 축소형 (한쪽에서 시작, 같은 방향)
한 포인터는 빠르게, 다른 포인터는 천천히 같은 방향으로 움직인다. "조건을 만족하는 가장 짧은/긴 구간"을 찾을 때 사용한다.
i = 0
j = 0
while i < len(arr):
# i로 윈도우 확장
window_size = i - j + 1
# 조건 확인 및 j 진행
while condition:
# j로 윈도우 축소
j += 1
i += 1
합이 s 이상인 가장 짧은 부분배열을 찾는 예시다. (예: 부분합)
arr = [2, 3, 1, 2, 4, 3]
s = 7
i = 0
j = 0
min_len = float('inf')
current_sum = 0
while i < len(arr):
current_sum += arr[i]
i += 1
while current_sum >= s:
min_len = min(min_len, i - j)
current_sum -= arr[j]
j += 1
print(min_len) # 2 ([4, 3])
양 끝 수렴형 (양 끝에서 시작, 좁혀옴)
정렬된 배열의 양 끝에서 시작해서 중간으로 좁혀온다. 두 원소의 합이 특정 조건을 만족하는 쌍을 찾을 때 사용한다.
left = 0
right = len(arr) - 1
while left < right:
total = arr[left] + arr[right]
if total > target:
right -= 1 # 합이 크면 오른쪽을 줄이고
elif total < target:
left += 1 # 합이 작으면 왼쪽을 늘린다
else:
break # 정확히 맞으면 종료
정렬된 배열에서 두 수의 합이 0에 가장 가까운 쌍을 찾는 예시다. (예: 두 용액)
arr = [-99, -2, -1, 4, 98]
left = 0
right = len(arr) - 1
best = float('inf')
ans_l, ans_r = 0, 0
while left < right:
total = arr[left] + arr[right]
if abs(total) < best:
best = abs(total)
ans_l, ans_r = arr[left], arr[right]
if total > 0:
right -= 1
elif total < 0:
left += 1
else:
break
print(ans_l, ans_r) # -99 98
연습 문제
9. 구간합 (Prefix Sum)
구간합은 미리 누적합 배열을 만들어두고, 구간의 합을 O(1)에 구하는 기법이다.
1D Prefix Sum
미리 합 배열을 만들어두면 구간 합을 뺄셈 한 번으로 구할 수 있다.
n = len(arr)
prefix = [0] * (n + 1)
for i in range(n):
prefix[i+1] = prefix[i] + arr[i]
def range_sum(l, r):
return prefix[r+1] - prefix[l]
1D prefix sum 사용 예시다. (예: 구간 합 구하기 4)
arr = [1, 2, 3, 4, 5]
n = len(arr)
prefix = [0] * (n + 1)
for i in range(n):
prefix[i+1] = prefix[i] + arr[i]
# [1, 3] 범위 합
print(prefix[4] - prefix[1]) # 9 (2+3+4)
2D Prefix Sum
2D에서는 직사각형 영역의 합을 미리 계산해둔다. 위 + 왼 - 대각 + 현재값으로 prefix 배열을 채운다.
rows, cols = len(grid), len(grid[0])
prefix = [[0] * (cols + 1) for _ in range(rows + 1)]
for i in range(1, rows + 1):
for j in range(1, cols + 1):
prefix[i][j] = grid[i-1][j-1] \
+ prefix[i-1][j] \
+ prefix[i][j-1] \
- prefix[i-1][j-1]
전체 - 위 - 왼 + 겹침으로 직사각형 구간 합을 구한다.
def range_sum_2d(r1, c1, r2, c2):
return prefix[r2+1][c2+1] \
- prefix[r1][c2+1] \
- prefix[r2+1][c1] \
+ prefix[r1][c1]
2D prefix sum 전체 코드다. (예: 구간 합 구하기 5)
grid = [[1, 2], [3, 4], [5, 6]]
rows, cols = len(grid), len(grid[0])
prefix = [[0] * (cols + 1) for _ in range(rows + 1)]
for i in range(1, rows + 1):
for j in range(1, cols + 1):
prefix[i][j] = grid[i-1][j-1] \
+ prefix[i-1][j] \
+ prefix[i][j-1] \
- prefix[i-1][j-1]
def range_sum_2d(r1, c1, r2, c2):
return prefix[r2+1][c2+1] \
- prefix[r1][c2+1] \
- prefix[r2+1][c1] \
+ prefix[r1][c1]
print(range_sum_2d(0, 0, 2, 1)) # 1+2+3+4+5+6 = 21
정수로 할 수 있는 계산은 먼저 정수로 하고, 나누기는 마지막에 한다.
# 비교: a/b < c/d
# X: a/b < c/d (부동소수점 오차)
# O: a*d < c*b (정수 연산)
연습 문제
10. BFS (너비 우선 탐색)
기본 템플릿
큐에 시작점을 넣고, 인접한 노드를 방문하며 탐색한다. (예: DFS와 BFS) visited로 중복 방문을 막는다.
from collections import deque
def bfs(start):
queue = deque([start])
visited = {start}
while queue:
node = queue.popleft()
# node 처리
for neighbor in graph[node]:
if neighbor not in visited:
visited.add(neighbor)
queue.append(neighbor)
set 방식 (자주 사용)
visited를 set으로 관리하는 방식이다. 노드가 좌표나 문자열일 때 유용하다.
from collections import deque
def bfs(start, graph):
queue = deque([start])
visited = {start}
while queue:
node = queue.popleft()
for neighbor in graph[node]:
if neighbor not in visited:
visited.add(neighbor)
queue.append(neighbor)
return visited
멀티소스 BFS
여러 개의 시작점에서 동시에 BFS를 수행한다. 시작점을 전부 큐에 넣고 시작하면 각 노드까지의 최단 거리를 구할 수 있다. (예: 토마토)
from collections import deque
def multi_source_bfs(sources, graph):
queue = deque()
dist = {source: 0 for source in sources}
for source in sources:
queue.append(source)
while queue:
node = queue.popleft()
for neighbor in graph[node]:
if neighbor not in dist:
dist[neighbor] = dist[node] + 1
queue.append(neighbor)
return dist
간선 비용이 모두 1인 그래프에서 멀티소스 BFS를 수행하는 예시다.
from collections import deque
def shortest_distance(n, edges, sources):
# 그래프 구성
graph = [[] for _ in range(n)]
for u, v in edges:
graph[u].append(v)
graph[v].append(u)
# 멀티소스 BFS
queue = deque()
dist = [-1] * n
for source in sources:
dist[source] = 0
queue.append(source)
while queue:
node = queue.popleft()
for neighbor in graph[node]:
if dist[neighbor] == -1:
dist[neighbor] = dist[node] + 1
queue.append(neighbor)
return dist
격자 BFS
2D 격자에서 상하좌우로 이동하며 BFS를 수행한다. (예: 미로 탐색)
from collections import deque
def grid_bfs(grid, start_r, start_c):
rows, cols = len(grid), len(grid[0])
queue = deque([(start_r, start_c)])
visited = {(start_r, start_c)}
# 상하좌우 이동
dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]
while queue:
x, y = queue.popleft()
for i in range(4):
nx, ny = x + dx[i], y + dy[i]
if 0 <= nx < rows and 0 <= ny < cols \
and (nx, ny) not in visited \
and grid[nx][ny] != 0: # 이동 가능한 칸
visited.add((nx, ny))
queue.append((nx, ny))
return visited
연습 문제
11. DFS (깊이 우선 탐색)
재귀 방식
함수가 자기 자신을 호출하며 깊이 탐색한다. (예: DFS와 BFS) 코드가 간결하지만 재귀 깊이 제한에 주의한다.
def dfs_recursive(node, graph, visited):
visited.add(node)
for neighbor in graph[node]:
if neighbor not in visited:
dfs_recursive(neighbor, graph, visited)
스택 방식
재귀 대신 스택을 직접 사용한다. 재귀 깊이 제한을 피할 수 있다.
def dfs_stack(start, graph):
stack = [start]
visited = {start}
while stack:
node = stack.pop()
for neighbor in graph[node]:
if neighbor not in visited:
visited.add(neighbor)
stack.append(neighbor)
return visited
파이썬 재귀 깊이 기본값은 1000이다. 깊은 재귀가 필요하면 늘려야 한다.
# sys.setrecursionlimit 설정 (재귀 깊이 제한)
import sys
sys.setrecursionlimit(10**6)
12. 백트래킹
위 itertools 섹션에서는 라이브러리로 간단히 생성하는 방법을 다뤘다. itertools는 모든 경우를 전부 생성한 뒤 필터링하기 때문에, 예를 들어 '합이 15 이하인 조합만'과 같은 조건이 있으면 불필요한 경우까지 전부 만들어서 느리다. 백트래킹은 만드는 도중에 조건에 안 맞으면 중단하고 다른 경우로 넘어가기 때문에 훨씬 빠르다.
순열 (permutation)
순서를 고려하여 모든 경우를 직접 생성한다. 놓고 → 재귀 → 빼기(백트래킹)가 핵심이다.
def permute(arr, l, r, result):
if l == r:
result.append(arr[:])
else:
for i in range(l, r + 1):
arr[l], arr[i] = arr[i], arr[l]
permute(arr, l + 1, r, result)
arr[l], arr[i] = arr[i], arr[l]
arr = [1, 2, 3]
result = []
permute(arr, 0, len(arr) - 1, result)
# [[1,2,3], [1,3,2], [2,1,3], ...]
조합 (combination)
start 파라미터로 이전에 선택한 것보다 뒤의 것만 선택해서 중복을 방지한다.
def combine(arr, r, start, path, result):
if len(path) == r:
result.append(path[:])
else:
for i in range(start, len(arr)):
path.append(arr[i])
combine(arr, r, i + 1, path, result)
path.pop()
arr = [1, 2, 3]
result = []
combine(arr, 2, 0, [], result)
# [[1,2], [1,3], [2,3]]
중복순열 (repeated permutation)
used 배열 없이 모든 원소를 매번 선택할 수 있다.
def repeated_perm(arr, n, path, result):
if len(path) == n:
result.append(path[:])
else:
for item in arr:
path.append(item)
repeated_perm(arr, n, path, result)
path.pop()
arr = [1, 2]
result = []
repeated_perm(arr, 2, [], result)
# [[1,1], [1,2], [2,1], [2,2]]
중복조합 (repeated combination)
조합에서 같은 원소를 여러 번 선택할 수 있다. 다음 재귀에서 i + 1 대신 i부터 시작하면 된다.
def repeated_comb(arr, r, start, path, result):
if len(path) == r:
result.append(path[:])
else:
for i in range(start, len(arr)):
path.append(arr[i])
repeated_comb(arr, r, i, path, result) # i+1이 아니라 i
path.pop()
arr = [1, 2, 3]
result = []
repeated_comb(arr, 2, 0, [], result)
# [[1,1], [1,2], [1,3], [2,2], [2,3], [3,3]]
N-Queen (패턴)
queens 배열에 인덱스=행, 값=열로 저장한다. 각 행에 퀸이 하나뿐이므로 1D 배열로 충분하다.
def is_safe(queens, row, col):
# 같은 열 확인
for r in range(row):
if queens[r] == col:
return False
# 대각선 확인
for r in range(row):
if abs(r - row) == abs(queens[r] - col):
return False
return True
def solve(n, row, queens, result):
if row == n:
result.append(queens[:])
else:
for col in range(n):
if is_safe(queens, row, col):
queens[row] = col
solve(n, row + 1, queens, result)
queens[row] = -1
n = 4
result = []
solve(n, 0, [-1] * n, result)
# [[1, 3, 0, 2], [2, 0, 3, 1]]
연습 문제
13. 이진 탐색
기본 템플릿
정렬된 배열에서 target을 O(log n)에 찾는다. (예: 수 찾기) left와 right를 좁혀가며 mid를 비교한다.
def binary_search(arr, target):
left, right = 0, len(arr) - 1
while left <= right:
mid = (left + right) // 2
if arr[mid] == target:
return mid
elif arr[mid] < target:
left = mid + 1
else:
right = mid - 1
return -1
값이 있는 가장 왼쪽 위치
target 이상인 첫 번째 위치를 찾는다. bisect_left와 동일하다.
def binary_search_left(arr, target):
left, right = 0, len(arr)
while left < right:
mid = (left + right) // 2
if arr[mid] < target:
left = mid + 1
else:
right = mid
return left
값이 있는 가장 오른쪽 위치
target 이하인 마지막 위치를 찾는다. bisect_right에서 1을 뺀 값과 동일하다.
def binary_search_right(arr, target):
left, right = 0, len(arr)
while left < right:
mid = (left + right) // 2
if arr[mid] <= target:
left = mid + 1
else:
right = mid
return left - 1
답이 범위에 있는 경우
특정 조건을 만족하는 최소/최대값을 찾을 때 사용한다. 파라메트릭 서치라고도 한다.
def can_achieve(x):
# x로 조건을 만족할 수 있는지 확인
return True or False
left, right = 0, 10**9
while left < right:
mid = (left + right) // 2
if can_achieve(mid):
right = mid
else:
left = mid + 1
print(left) # 답
14. Union-Find (Disjoint Set Union)
서로소 집합을 관리하며 두 원소가 같은 집합에 속하는지 확인할 수 있다.
기본 구조
find로 루트를 찾고, union으로 두 집합을 합친다. (예: 집합의 표현) find를 호출할 때 거쳐간 노드들을 루트에 직접 연결하는 것을 경로 압축이라 하고, 다음 find가 O(1)에 가까워진다.
class UnionFind:
def __init__(self, n):
self.parent = list(range(n))
self.rank = [0] * n
def find(self, x):
if self.parent[x] != x:
self.parent[x] = self.find(self.parent[x]) # 경로 압축
return self.parent[x]
def union(self, x, y):
px, py = self.find(x), self.find(y)
if px == py:
return
# 랭크가 낮은 트리를 높은 트리 아래로
if self.rank[px] < self.rank[py]:
px, py = py, px
self.parent[py] = px
if self.rank[px] == self.rank[py]:
self.rank[px] += 1
def is_connected(self, x, y):
return self.find(x) == self.find(y)
사용 예제
union으로 합치고, is_connected로 같은 집합인지 확인한다.
uf = UnionFind(5)
uf.union(0, 1)
uf.union(2, 3)
uf.union(1, 2)
print(uf.is_connected(0, 3)) # True (연결됨)
print(uf.is_connected(0, 4)) # False
15. 그리디
그리디는 매 순간 최선의 선택을 반복해서 전체 최적해를 구하는 기법이다. 단, "매 순간 최선 = 전체 최선"이 성립하는 문제에서만 쓸 수 있다. 예를 들어 동전 거스름돈에서 큰 동전부터 쓰는 건 그리디가 맞지만, 배낭 문제에서 비싼 것부터 넣으면 최적이 아닐 수 있다.
정렬 후 처리 패턴
가장 큰(또는 작은) 것부터 처리하는 패턴이다. (예: 동전 0)
# 동전 거스름돈 (가장 큰 동전부터)
coins = [100, 50, 10, 1]
amount = 237
count = 0
for coin in coins:
count += amount // coin
amount %= coin
print(count) # 최소 동전 개수
구간 스케줄링 패턴
겹치지 않는 구간을 최대한 많이 선택하는 패턴이다. 끝나는 시간 기준으로 정렬한 뒤, 겹치지 않으면 선택한다. (예: 회의실 배정)
meetings = [(1, 3), (2, 5), (4, 6), (6, 7)]
# (시작 시간, 끝나는 시간)
meetings.sort(key=lambda x: x[1]) # 끝나는 시간 기준
count = 0
last_end = 0
for start, end in meetings:
if start >= last_end:
count += 1
last_end = end
print(count) # 최대 회의 개수
연습 문제
16. 동적 프로그래밍 (DP)
DP는 작은 문제의 답을 이용해서 큰 문제의 답을 구하는 기법이다. 반복되는 계산을 저장해서 재사용한다.
탑다운 vs 바텀업
DP를 구현하는 두 가지 방식이 있다.
바텀업 (Bottom-Up): 작은 문제부터 for문으로 채워 올라간다. 재귀 제한이 없고 직관적이다.
# 피보나치 - 바텀업
dp = [0] * (n + 1)
dp[1] = 1
for i in range(2, n + 1):
dp[i] = dp[i-1] + dp[i-2]
탑다운 (Top-Down): 큰 문제에서 시작해 재귀로 작은 문제를 호출한다. lru_cache나 memo 딕셔너리로 중복 계산을 방지한다.
# 피보나치 - 탑다운
from functools import lru_cache
@lru_cache(maxsize=None)
def fib(n):
if n < 2:
return n
return fib(n-1) + fib(n-2)
| 바텀업 | 탑다운 | |
|---|---|---|
| 구현 | for문 + dp 배열 | 재귀 + 메모이제이션 |
| 방향 | 작은 문제 → 큰 문제 | 큰 문제 → 작은 문제 |
| 장점 | 재귀 제한 없음 | 필요한 것만 계산 |
| 단점 | 불필요한 것도 계산할 수 있음 | 파이썬 재귀 깊이 제한 |
코테에서는 바텀업을 더 자주 사용한다. 재귀 제한 걱정이 없고 속도도 빠르다.
0/1 배낭 (Knapsack)
n개의 물건 중 무게 W를 초과하지 않으면서 최대 가치를 담는 문제다. (예: 평범한 배낭)
dp[i][j]는 처음 i개 물건 중에서 무게 j 이하로 담을 때의 최대 가치다. 각 물건마다 "안 담기 vs 담기" 중 큰 값을 선택한다.
def knapsack(n, W, weights, values):
dp = [[0] * (W + 1) for _ in range(n + 1)]
for i in range(1, n + 1):
w, v = weights[i-1], values[i-1]
for j in range(W + 1):
dp[i][j] = dp[i-1][j] # 안 담기
if j >= w:
dp[i][j] = max(dp[i][j], dp[i-1][j-w] + v) # 담기
return dp[n][W]
2D 배열 대신 1D 배열로 공간을 절약할 수 있다. 역순으로 순회해야 같은 물건을 중복으로 담는 것을 방지한다.
def knapsack_1d(n, W, weights, values):
dp = [0] * (W + 1)
for i in range(n):
w, v = weights[i], values[i]
for j in range(W, w - 1, -1): # 역순 순회
dp[j] = max(dp[j], dp[j-w] + v)
return dp[W]
연습 문제
17. 세그먼트 트리
배열의 구간 합을 빠르게 구하고, 값을 빠르게 변경할 수 있는 자료구조다.
| 연산 | Prefix Sum | 세그먼트 트리 |
|---|---|---|
| 구간합 구하기 | O(1) | O(log n) |
| 값 변경 | O(n) | O(log n) |
값이 자주 변하면서 구간합도 자주 구해야 할 때 세그먼트 트리를 사용한다.
개념
배열을 완전 이진 트리로 표현한다. 리프 노드는 원소, 내부 노드는 자식들의 합을 저장한다.
배열을 완전 이진 트리로 표현한다. tree[i]의 왼쪽 자식은 tree[2*i], 오른쪽 자식은 tree[2*i+1]이다. 리프 노드에는 원본 배열의 원소가 들어가고, 내부 노드에는 자식들의 합이 저장된다.
전체 템플릿
build로 초기 트리를 구성하고 (예: 구간 합 구하기), update로 값을 수정하며, query로 구간 합을 구한다.
tree = [0] * (4 * n)
def build(node, start, end):
if start == end:
tree[node] = arr[start]
else:
mid = (start + end) // 2
build(2 * node, start, mid)
build(2 * node + 1, mid + 1, end)
tree[node] = tree[2 * node] + tree[2 * node + 1]
def update(node, start, end, idx, val):
if start == end:
tree[node] = val
else:
mid = (start + end) // 2
if idx <= mid:
update(2 * node, start, mid, idx, val)
else:
update(2 * node + 1, mid + 1, end, idx, val)
tree[node] = tree[2 * node] + tree[2 * node + 1]
def query(node, start, end, l, r):
if r < start or end < l:
return 0
if l <= start and end <= r:
return tree[node]
mid = (start + end) // 2
return query(2 * node, start, mid, l, r) + query(2 * node + 1, mid + 1, end, l, r)
# 사용
build(1, 0, n - 1) # 트리 생성
update(1, 0, n - 1, idx, val) # 값 변경
print(query(1, 0, n - 1, l, r)) # 구간 합
연습 문제
부록 1. 시간초과 대처
sys.stdin.readline 사용
input()보다 빠르다. 반복 입력이 많을 때 사용한다.
import sys
input = sys.stdin.readline
n = int(input())
for _ in range(n):
line = input().split()
sys.setrecursionlimit
파이썬의 기본 재귀 제한은 1000이다. 깊은 재귀가 필요하면 늘려야 한다.
import sys
sys.setrecursionlimit(10**6)
PyPy3로 제출
같은 코드를 PyPy3로 제출하면 보통 2-3배 빠르다. 대부분의 온라인 저지에서 지원한다.
print 최적화
print()를 반복 호출하면 느리다. join으로 모아서 한 번에 출력한다.
# X: 느림
for i in range(10000):
print(i)
# O: 빠름
result = [str(i) for i in range(10000)]
print('\n'.join(result))
부록 2. 부동소수점 주의
정수 우선 계산
비교 연산은 나누기 대신 곱셈으로 변환해서 부동소수점 오차를 피한다.
# X: 부동소수점 오차
if a / b < c / d:
pass
# O: 정수 연산
if a * d < c * b:
pass
나누기는 마지막에
정수로 할 수 있는 계산은 먼저 정수로 처리하고, 나누기는 마지막 단계에서만 수행한다.
# 분산 계산
# X: (1/n * sum(x^2)) - (1/n * sum(x))^2
# O: (sum(x^2) / n) - (sum(x) / n)^2
# 더 좋음: (n * sum(x^2) - (sum(x))^2) / n^2
n = 100
sum_x = 5000
sum_x2 = 250000
variance = (n * sum_x2 - sum_x * sum_x) / (n * n)
부록 3. 파일 입출력 (Output-Only 문제)
일부 문제는 표준 입출력 대신 파일 입출력을 요구한다.
파일 읽기/쓰기
with문으로 파일을 열고, readlines()로 전체를 읽거나 write()로 결과를 쓴다.
# 읽기
with open('input.txt', 'r') as f:
lines = f.readlines()
for line in lines:
line = line.strip()
# 또는
with open('input.txt', 'r') as f:
data = f.read()
# 쓰기
with open('output.txt', 'w') as f:
f.write('result\n')
f.write('answer\n')
한 줄씩 읽기
readline()을 반복 호출해서 한 줄씩 처리한다.
with open('input.txt', 'r') as f:
n = int(f.readline())
for _ in range(n):
line = f.readline().strip()
# 처리
부록 4. 자주 하는 실수
1. 리스트 초기화 실수
[[0]*m]*n은 모든 행이 같은 객체를 참조해서 한 행을 수정하면 나머지 행도 바뀐다.
# X: 모든 행이 같은 객체를 참조
grid = [[0] * 5] * 3
grid[0][0] = 1
print(grid[1][0]) # 1 (예상: 0)
# O
grid = [[0] * 5 for _ in range(3)]
2. sort() 체이닝
sort()는 None을 반환하므로 체이닝하면 에러가 난다. sorted()는 새 리스트를 반환한다.
# X
arr.sort().sort() # None 반환
# O
arr.sort()
arr.sort()
3. divmod 활용 안 함
몫과 나머지를 동시에 구할 때는 divmod()로 한 줄로 쓴다.
# X
quotient = a // b
remainder = a % b
# O
quotient, remainder = divmod(a, b)
4. 무한 루프 주의
이미 방문한 노드를 다시 큐에 넣으면 무한 루프가 발생한다. visited 체크를 먼저 해야 한다.
# BFS에서
while queue:
# queue에 계속 추가하면서 처리
# 주의: 이미 방문한 노드 다시 추가하지 않기
if neighbor not in visited:
visited.add(neighbor)
queue.append(neighbor)
5. 백트래킹에서 배열 복사
가변 리스트를 그대로 저장하면 이후 변경이 반영된다. [:]나 list()로 복사해서 저장한다.
# X: 모두 같은 리스트 참조
result = []
for perm in permutations:
result.append(perm) # perm이 변경되면 result의 모든 원소가 변함
# O: 명시적으로 복사
result.append(perm[:]) # 또는 list(perm)