코딩 테스트를 풀다 보면 알고리즘 자체보다 기본 문법이나 표준 라이브러리가 손에 익지 않아서 시간이 걸릴 때가 있다. 자주 쓰는 리스트, 문자열, 딕셔너리 패턴은 따로 외운다기보다 여러 문제를 풀면서 바로 떠올릴 수 있을 정도로 익혀두는 편이 좋다.

입력은 sys.stdin으로

간단한 문제에서는 input()이 편하지만, 입력이 수만 줄이 되면 input()만으로 시간 초과가 나기도 한다. 입력이 많은 문제는 sys.stdin.readline을 쓰는 것이 기본이다.

import sys
input = sys.stdin.readline
 
n = int(input())
arr = list(map(int, input().split()))

readline()은 줄 끝의 개행 문자를 포함해서 돌려주므로, 문자열로 쓸 때는 .rstrip()으로 잘라내야 한다. 숫자로 변환할 때는 int()가 공백과 개행을 알아서 무시하니 신경 쓰지 않아도 된다.

리스트

리스트는 코딩 테스트의 기본 자료구조다. 생성, 추가, 정렬, 뒤집기가 기본 패턴이다.

arr = [3, 1, 4, 1, 5]
 
arr.append(9)        # 뒤에 추가 — O(1)
arr.sort()           # 제자리 정렬 — O(n log n)
arr.reverse()        # 제자리 뒤집기
top = arr[-1]        # 마지막 원소
sub = arr[1:4]       # 슬라이싱 — 새 리스트 O(k)
 
grid = [[0] * m for _ in range(n)]   # 2차원 배열

2차원 배열을 만들 때 [[0] * m] * n으로 쓰면 안 된다. 같은 행 객체가 n번 복사되어서, 한 칸을 바꾸면 모든 행이 같이 바뀐다. 반드시 리스트 컴프리헨션으로 행마다 새로 만들어야 한다.

문자열

문자열에서는 분리, 결합, 변환을 많이 쓴다.

s = "algorithm test"
 
words = s.split()          # ['algorithm', 'test']
joined = "-".join(words)   # 'algorithm-test'
up = s.upper()
found = s.find("test")     # 없으면 -1 (index()는 예외)
num = ord('a')             # 97 — 문자 ↔ 코드
ch = chr(97)               # 'a'

파이썬 문자열은 불변이라 반복문 안에서 s += ch로 이어 붙이면 매번 새 문자열이 생긴다. 조각을 리스트에 모았다가 마지막에 "".join(pieces)로 합치는 것이 정석이다.

딕셔너리와 집합

“존재하는지”, “몇 번 나왔는지”를 물으면 딕셔너리와 집합부터 떠올린다. 리스트의 inO(n)이지만 집합과 딕셔너리의 inO(1)이다.

from collections import Counter, defaultdict
 
seen = set()
seen.add(3)
if 3 in seen: ...          # O(1)
 
count = Counter("aabbbc")  # {'b': 3, 'a': 2, 'c': 1}
count.most_common(1)       # [('b', 3)]
 
graph = defaultdict(list)  # 키가 없어도 기본값 생성
graph[1].append(2)

정렬 키와 람다

sort()sorted()key 인자는 문제 조건을 그대로 옮기는 자리다. “길이 순, 같으면 사전순” 같은 다중 조건은 튜플로 표현한다.

words.sort(key=lambda w: (len(w), w))       # 길이 → 사전순
pairs.sort(key=lambda p: -p[1])             # 두 번째 값 내림차순

함정

  • list.pop(0)O(n)이다. 큐가 필요하면 collections.dequepopleft()를 쓴다.
  • 재귀 한도. 파이썬 기본 재귀 한도는 1,000이라 DFS에서 바로 터진다. sys.setrecursionlimit(10**6)을 습관처럼 올려두거나 스택 반복문으로 바꾼다.
  • 얕은 복사. b = a는 같은 리스트를 가리킨다. 복사가 필요하면 a[:] 또는 a.copy(), 2차원이면 copy.deepcopy.
  • 부동소수점 비교. 나눗셈 결과를 ==로 비교하지 않는다. 가능하면 곱셈으로 바꿔 정수로 비교한다.