코딩테스트 준비 collections, heapq 활용

collections, heapq를 사용해보자

파이썬을 이용한 코딩테스트 연습시 원하고자 하는 기능을 리스트와 딕셔너리로 구현할때 시간초과나 메모리초과라는 문제점을 만날때가 있다.

1. 카운팅 기능을 리스트에서 구현할때

count 기능을 for문 사용시

#count
words = [0,1,2,3,4,2,3,1,2,4,0]
for i in range(len(words)):
    print(words.count(i),end=" ")
# 출력결과 > 1 2 3 2

이때의 시간복잡도는 O(N)인데 for문 안에서 n번 반복하기 때문에 O(N^2)이 됩니다.

하지만 collections의 Counter를 사용시

#Counter
from collections import Counter
words = [0,1,2,3,4,2,3,1,2,4,0]
counter = Counter(words)
for i in range(len(words)):
    print(counter[i],end=" ")

시간복잡도는 O(1) for문을 통해 n번 반복하기 때문에 O(N)이 된다.

또한 딕셔너리와 같은 기능을 제공하기 때문에 앞으로 카운팅 기능이 필요할때 딕셔너리에 저장하는 것보다 그냥 collections의 Counter기능을 쓸때 시간복잡도를 더 줄일 수 있다.

2. 리스트의 공간복잡도가 커진다면?

queue 기능을 파이썬에서 구현하고자 한다면 heapq를 사용하면 좋다 예를 들어 리스트에서 sort시켜 맨 앞의 작은 수끼리 더해나가는 자료구조를 생각할때 리스트로 구현하면 메모리초과가 날 수 있다.

# 1715, 카드 정렬하기
import sys

n = int(sys.stdin.readline())
ns = list(map(int,sys.stdin.read().split()))
ns.sort()
temp = 0
for i in range(len(ns)-1):
    temp += ns[i] + ns[i+1]
    ns[i+1] = temp
print(temp)

# 1715, 카드 정렬하기
import sys
import heapq


n = int(sys.stdin.readline())
ns = list(map(int,sys.stdin.read().split()))
if len(ns) == 1:
    print(0)
else:
    ans = 0
    while len(ns) > 1:
        temp1 = heapq.heappop(ns)
        temp2 = heapq.heappop(ns)
        ans += temp1 + temp2
        heapq.heappush(ns, temp1 + temp2)
    print(ans)

리스트를 통해 구현할때보다 메모리를 아낄 수 있다.