시작

파이썬으로 PS할 때에 유용한 팁들을 정리해봤다.

lru_cache

함수의 반환값을 기억해주는 lru_cache

2748번: 피보나치 수 2

def fibo(n):
    if n < 2:
        return n
    return fibo(n - 1) + fibo(n - 2)

print(fibo(int(input())))

이런 DP문제는 위처럼 단순 재귀로 풀면 당연히 TLE가 뜬다.

그래서 메모이제이션(memoization)을 하면서 반복적인 계산을 줄여줘야 한다.

그런데 이 부분을 정말 쉽게 해결해주는 게 바로 functools모듈의 lru_cache함수다.

from functools import lru_cache

@lru_cache(maxsize=None)
def fibo(n):
    if n < 2:
        return n
    return fibo(n - 1) + fibo(n - 2)

print(fibo(int(input())))

lru_cache는 데코레이터로 사용되는데, 원하는 함수에 이 데코레이터를 달아 주면 반환값을 알아서 저장해준다.

전자가 lru_cache를 사용한 코드, 후자가 데코레이터가 없는 코드다.

readline

반복해서 입력받아야 한다면 sys모듈의 stdin.readline을 사용하자.

import sys

n = sys.stdin.readline()

단순히 한두번 입력받고 마는 문제라면 괜찮지만, 입력을 많이 받을 경우에 기본 input함수를 사용하면 상당히 느려진다.

예시로 1725번: 히스토그램과 같이 여러번 입력받아야 하는 문제의 경우, readline을 사용하면 시간이 엄청나게 단축되는 것을 볼 수 있다.

전자가 readline사용, 후자가 기본 input을 사용했을 때의 시간이다.

from sys import stdin

input = stdin.readline

나는 항상 readline을 사용하려고 편하게 코드 맨 위에 이 코드를 적어놓는다.

🚨 주의사항

int(stdin.readline())으로 사용할 때는 괜찮지만, 문자열을 입력받을 경우 맨 뒤의 개행 문자 (\n)까지 그대로 입력받게 된다.
그러니 반드시 stdin.readline().strip()과 같이 사용하여 개행 문자를 제거해주자.

Asterisk, sep in print

리스트의 요소들을 출력해야 할 때, Asterisk(*)와 sep인자를 잘 활용하면 간단하게 원하는 방식으로 출력할 수 있다.

l = [1, 2, 3, 4, 5]
for i in l:
    print(i, end=" ")

# 1 2 3 4 5


for i in l:
    print(i)

# 1
# 2
# 3
# 4
# 5

이번에 설명할 방법을 모른다면 위와 같은 방식으로 출력할 것이다.

하지만 *연산자를 사용하면, 첫 번째 출력은 다음과 같이 줄일 수 있다.

print(*l)

아무런 기타 구문 없이 한 줄만으로 출력되는것을 볼 수 있다.

그리고 두 번째 출력은 *연산자와 sep인자를 사용해서 다음과 같이 줄일 수 있다.

print(*l, sep="\n")

*인자를 단독으로 사용하면 요소들 사이에 공백이 들어가는데, sep인자를 사용하면 요소들 사이에 들어갈 문자를 바꿔줄 수 있다.

위 예제에서는 요소들 사이에 들어갈 문자를 개행문자로 바꿔 한 줄에 하나씩 출력하게 된다.

map

map은 이터레이터(iterator)이다.

입력을 받을 때 map함수를 사용하는 경우가 많을 것이다.

m = map(int, input().split())

만약 1 2 3 4 5와 같은 입력값을 map을 이용해 위와 같이 입력받았는데, for문으로 하나씩 사용해야 한다면 어떻게 하겠는가?

l = list(m)
for i in l:
    print(i)

# 1
# ...

생각보다 위처럼 list로 바꿔주는 사람들이 많았다. 그런데 사실 map객체도 반복 가능한 객체(iterator)이다.

for i in m:
    print(i)

# 1
# ...

그래서 그냥 위처럼만 해 줘도 잘 돌아가는 것을 볼 수 있다.

print(*m)
# 1 2 3 4 5

마찬가지로, 위에서 설명한 print의 *연산자를 활용한 예제에도 사용할 수 있다.

join

문자열 결합은 join함수를 사용하자.

join함수는 문자열들로 구성된 iterator의 요소들을 결합해주는 함수이다.

l = ["a", "b", "c"]

", ".join(l)
# a, b, c

"".join(l)
# abc

정말 편리하게 요소들 사이에만 원하는 문자열을 추가해줘서, 맨 처음이나 마지막에 따로 처리를 할 필요가 없어진다.

set

중복 제거가 필요할 때는 set

set객체는 대부분 알지만 많이 사용해보진 않았을 것이다.

>>> l = [1, 2, 5, 3, 4, 1]
>>> set(l)
{1, 2, 3, 4, 5}

set에서는 중복을 제거해줄 뿐만 아니라 무려 정렬도 해 준다.

setrecursionlimit

파이썬에서 재귀함수를 사용한다면 setrecursionlimit는 반 필수

이름 그대로 재귀 깊이를 설정해주는 함수다.

파이썬에서는 기본 재귀 깊이가 1000밖에 되지 않기에 이 함수를 사용해서 깊이를 설정하고 문제를 푸는 게 좋다.

from sys import setrecursionlimit

setrecursionlimit(10 ** 4)

보통은 위와 같이 사용한다.

collections

적어도 Counter는 사용하는 법을 알아 두자.

파이썬의 내장 모듈 중 collections는 정말 유용한 모듈이다.

이 모듈에서 대표적인 객체는 deque, Counter등이 있는데,

deque은 해당 자료구조가 필요한 문제라고 암시하는 경우가 많지만 Counter는 정말 있으면 편한 순간이 많을 것이다.

from collections import Counter

c = Counter([1, 2, 3, 4, 5, 5, 5, 4, 2])
print(c)

# Counter({5: 3, 2: 2, 4: 2, 1: 1, 3: 1})

Counter는 iterator의 요소 개수를 세어 주는 객체라고 생각하면 된다.

게다가 c.most_common(), for element, index in list(c.items())등과 같은 정말 많은 함수, 사용 방법들이 있다.

여기서 설명하기에는 너무 많으니 직접 찾아보는걸 추천.

시간복잡도 관련

in

객체 안에 요소가 있는지 확인하는 in연산자는 list에서 사용할 경우 $O(N)$의 시간복잡도를 가지지만, set에서 사용할 경우 무려 $O(1)$의 시간복잡도를 갖는다.

math.gcd의 TMI

파이썬의 내장모듈 math에 들어있는 최대공약수를 구하는 함수 gcd는 유클리드 알고리즘을 사용하며, $O(logN)$의 시간복잡도를 가진다.

끝

파이썬은 신이야